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 credential_status<'a>(
76        &'a self,
77        subject: &'a str,
78    ) -> BoxFuture<'a, Result<Option<UpstreamOAuthCredentialStatus>, UpstreamOAuthError>>;
79
80    fn clear_credentials<'a>(
81        &'a self,
82        subject: &'a str,
83    ) -> BoxFuture<'a, Result<(), UpstreamOAuthError>>;
84
85    fn access_token<'a>(
86        &'a self,
87        subject: &'a str,
88    ) -> BoxFuture<'a, Result<String, UpstreamOAuthError>>;
89}
90
91#[derive(Clone)]
92pub struct UpstreamOAuthRuntime {
93    provider: Arc<dyn UpstreamOAuthProvider>,
94    managers: Arc<BTreeMap<String, Arc<dyn UpstreamOAuthManager>>>,
95}
96
97impl UpstreamOAuthRuntime {
98    #[must_use]
99    pub fn new(
100        provider: Arc<dyn UpstreamOAuthProvider>,
101        managers: BTreeMap<String, Arc<dyn UpstreamOAuthManager>>,
102    ) -> Self {
103        Self {
104            provider,
105            managers: Arc::new(managers),
106        }
107    }
108
109    #[must_use]
110    pub fn provider(&self) -> Arc<dyn UpstreamOAuthProvider> {
111        Arc::clone(&self.provider)
112    }
113
114    #[must_use]
115    pub fn manager(&self, upstream: &str) -> Option<Arc<dyn UpstreamOAuthManager>> {
116        self.managers.get(upstream).cloned()
117    }
118
119    pub fn evict_subject(&self, upstream: &str, subject: &str) {
120        self.provider.evict_subject(upstream, subject);
121    }
122
123    pub fn evict_upstream(&self, upstream: &str) {
124        self.provider.evict_upstream(upstream);
125    }
126}
127
128#[cfg(test)]
129#[path = "oauth_tests.rs"]
130mod tests;
soma_mcp_client/oauth.rs

soma_mcp_client/
oauth.rs
