fix: Implement a CredentialStore for auth (#5741)

This commit is contained in:
Alex Hancock
2025-11-17 13:05:40 -05:00
committed by GitHub
parent 655d5d5b14
commit 2d7c7c8e07
12 changed files with 129 additions and 91 deletions
Generated
+44 -8
View File
@@ -281,9 +281,9 @@ dependencies = [
[[package]] [[package]]
name = "async-trait" name = "async-trait"
version = "0.1.87" version = "0.1.89"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d556ec1359574147ec0c4fc5eb525f3f23263a592b1a9c07e0a75b427de55c97" checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
@@ -2661,7 +2661,7 @@ dependencies = [
"rand 0.8.5", "rand 0.8.5",
"regex", "regex",
"reqwest 0.12.12", "reqwest 0.12.12",
"rmcp", "rmcp 0.9.0",
"schemars", "schemars",
"serde", "serde",
"serde_json", "serde_json",
@@ -2711,7 +2711,7 @@ dependencies = [
"once_cell", "once_cell",
"paste", "paste",
"regex", "regex",
"rmcp", "rmcp 0.9.0",
"serde", "serde",
"serde_json", "serde_json",
"tokio", "tokio",
@@ -2750,7 +2750,7 @@ dependencies = [
"open", "open",
"rand 0.8.5", "rand 0.8.5",
"regex", "regex",
"rmcp", "rmcp 0.9.0",
"rustyline", "rustyline",
"serde", "serde",
"serde_json", "serde_json",
@@ -2805,7 +2805,7 @@ dependencies = [
"rayon", "rayon",
"regex", "regex",
"reqwest 0.11.27", "reqwest 0.11.27",
"rmcp", "rmcp 0.8.5",
"schemars", "schemars",
"serde", "serde",
"serde_json", "serde_json",
@@ -2856,7 +2856,7 @@ dependencies = [
"goose-mcp", "goose-mcp",
"http 1.2.0", "http 1.2.0",
"reqwest 0.12.12", "reqwest 0.12.12",
"rmcp", "rmcp 0.9.0",
"schemars", "schemars",
"serde", "serde",
"serde_json", "serde_json",
@@ -5574,6 +5574,29 @@ version = "0.8.5"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e5947688160b56fb6c827e3c20a72c90392a1d7e9dec74749197aa1780ac42ca" checksum = "e5947688160b56fb6c827e3c20a72c90392a1d7e9dec74749197aa1780ac42ca"
dependencies = [ dependencies = [
"base64 0.22.1",
"chrono",
"futures",
"paste",
"pin-project-lite",
"rmcp-macros 0.8.5",
"schemars",
"serde",
"serde_json",
"thiserror 2.0.12",
"tokio",
"tokio-stream",
"tokio-util",
"tracing",
]
[[package]]
name = "rmcp"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "acc36ea743d4bbc97e9f3c33bf0b97765a5cf338de3d9c3d2f321a6e38095615"
dependencies = [
"async-trait",
"base64 0.22.1", "base64 0.22.1",
"chrono", "chrono",
"futures", "futures",
@@ -5583,7 +5606,7 @@ dependencies = [
"pin-project-lite", "pin-project-lite",
"process-wrap", "process-wrap",
"reqwest 0.12.12", "reqwest 0.12.12",
"rmcp-macros", "rmcp-macros 0.9.0",
"schemars", "schemars",
"serde", "serde",
"serde_json", "serde_json",
@@ -5609,6 +5632,19 @@ dependencies = [
"syn 2.0.99", "syn 2.0.99",
] ]
[[package]]
name = "rmcp-macros"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "263caba1c96f2941efca0fdcd97b03f42bcde52d2347d05e5d77c93ab18c5b58"
dependencies = [
"darling 0.21.0",
"proc-macro2",
"quote",
"serde_json",
"syn 2.0.99",
]
[[package]] [[package]]
name = "ron" name = "ron"
version = "0.8.1" version = "0.8.1"
+1 -1
View File
@@ -15,7 +15,7 @@ uninlined_format_args = "allow"
string_slice = "warn" string_slice = "warn"
[workspace.dependencies] [workspace.dependencies]
rmcp = { version = "0.8.5", features = ["schemars", "auth"] } rmcp = { version = "0.9.0", features = ["schemars", "auth"] }
# Patch for Windows cross-compilation issue with crunchy # Patch for Windows cross-compilation issue with crunchy
[patch.crates-io] [patch.crates-io]
+1 -1
View File
@@ -16,7 +16,7 @@ paste = "1.0"
ctor = "0.2.7" ctor = "0.2.7"
goose = { path = "../goose" } goose = { path = "../goose" }
rmcp = { workspace = true } rmcp = { workspace = true }
async-trait = "0.1.86" async-trait = "0.1.89"
chrono = { version = "0.4", features = ["serde"] } chrono = { version = "0.4", features = ["serde"] }
serde_json = "1.0" serde_json = "1.0"
serde = { version = "1.0", features = ["derive"] } serde = { version = "1.0", features = ["derive"] }
+1 -1
View File
@@ -43,7 +43,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt", "json",
tracing-appender = "0.2" tracing-appender = "0.2"
once_cell = "1.20.2" once_cell = "1.20.2"
shlex = "1.3.0" shlex = "1.3.0"
async-trait = "0.1.86" async-trait = "0.1.89"
base64 = "0.22.1" base64 = "0.22.1"
regex = "1.11.1" regex = "1.11.1"
nix = { version = "0.30.1", features = ["process", "signal"] } nix = { version = "0.30.1", features = ["process", "signal"] }
+1 -1
View File
@@ -33,7 +33,7 @@ reqwest = { version = "0.11", features = [
"json", "json",
"rustls-tls-native-roots", "rustls-tls-native-roots",
], default-features = false } ], default-features = false }
async-trait = "0.1" async-trait = "0.1.89"
chrono = { version = "0.4.38", features = ["serde"] } chrono = { version = "0.4.38", features = ["serde"] }
etcetera = "0.8.0" etcetera = "0.8.0"
tempfile = "3.8" tempfile = "3.8"
+1 -1
View File
@@ -50,5 +50,5 @@ path = "src/bin/generate_schema.rs"
[dev-dependencies] [dev-dependencies]
tower = "0.5" tower = "0.5"
async-trait = "0.1" async-trait = "0.1.89"
tempfile = "3.15.0" tempfile = "3.15.0"
+1 -1
View File
@@ -49,7 +49,7 @@ serde_urlencoded = "0.7"
jsonschema = "0.30.0" jsonschema = "0.30.0"
uuid = { version = "1.0", features = ["v4"] } uuid = { version = "1.0", features = ["v4"] }
regex = "1.11.1" regex = "1.11.1"
async-trait = "0.1" async-trait = "0.1.89"
async-stream = "0.3" async-stream = "0.3"
minijinja = { version = "2.10.2", features = ["loader"] } minijinja = { version = "2.10.2", features = ["loader"] }
include_dir = "0.7.4" include_dir = "0.7.4"
@@ -674,6 +674,7 @@ impl ExtensionManager {
output_schema: tool.output_schema, output_schema: tool.output_schema,
icons: None, icons: None,
title: None, title: None,
meta: None,
}); });
} }
} }
+25 -14
View File
@@ -1,9 +1,11 @@
mod persist;
use axum::extract::{Query, State}; use axum::extract::{Query, State};
use axum::response::Html; use axum::response::Html;
use axum::routing::get; use axum::routing::get;
use axum::Router; use axum::Router;
use minijinja::render; use minijinja::render;
use rmcp::transport::auth::OAuthState; use rmcp::transport::auth::{CredentialStore, OAuthState, StoredCredentials};
use rmcp::transport::AuthorizationManager; use rmcp::transport::AuthorizationManager;
use serde::Deserialize; use serde::Deserialize;
use std::net::SocketAddr; use std::net::SocketAddr;
@@ -11,9 +13,7 @@ use std::sync::Arc;
use tokio::sync::{oneshot, Mutex}; use tokio::sync::{oneshot, Mutex};
use tracing::warn; use tracing::warn;
use crate::oauth::persist::{clear_credentials, load_cached_state, save_credentials}; use crate::oauth::persist::GooseCredentialStore;
mod persist;
const CALLBACK_TEMPLATE: &str = include_str!("oauth_callback.html"); const CALLBACK_TEMPLATE: &str = include_str!("oauth_callback.html");
@@ -32,18 +32,21 @@ pub async fn oauth_flow(
mcp_server_url: &String, mcp_server_url: &String,
name: &String, name: &String,
) -> Result<AuthorizationManager, anyhow::Error> { ) -> Result<AuthorizationManager, anyhow::Error> {
if let Ok(oauth_state) = load_cached_state(mcp_server_url, name).await { let credential_store = GooseCredentialStore::new(name.clone());
if let Some(authorization_manager) = oauth_state.into_authorization_manager() { let mut auth_manager = AuthorizationManager::new(mcp_server_url).await?;
if authorization_manager.refresh_token().await.is_ok() { auth_manager.set_credential_store(credential_store.clone());
return Ok(authorization_manager);
} if auth_manager.initialize_from_store().await? {
if auth_manager.refresh_token().await.is_ok() {
return Ok(auth_manager);
} }
if let Err(e) = clear_credentials(name) { if let Err(e) = credential_store.clear().await {
warn!("error clearing bad credentials: {}", e); warn!("error clearing bad credentials: {}", e);
} }
} }
// No existing credentials or they were invalid - need to do the full oauth flow
let (code_sender, code_receiver) = oneshot::channel::<CallbackParams>(); let (code_sender, code_receiver) = oneshot::channel::<CallbackParams>();
let app_state = AppState { let app_state = AppState {
code_receiver: Arc::new(Mutex::new(Some(code_sender))), code_receiver: Arc::new(Mutex::new(Some(code_sender))),
@@ -74,6 +77,7 @@ pub async fn oauth_flow(
}); });
let mut oauth_state = OAuthState::new(mcp_server_url, None).await?; let mut oauth_state = OAuthState::new(mcp_server_url, None).await?;
let redirect_uri = format!("http://localhost:{}/oauth_callback", used_addr.port()); let redirect_uri = format!("http://localhost:{}/oauth_callback", used_addr.port());
oauth_state oauth_state
.start_authorization(&[], redirect_uri.as_str(), Some("goose")) .start_authorization(&[], redirect_uri.as_str(), Some("goose"))
@@ -91,13 +95,20 @@ pub async fn oauth_flow(
} = code_receiver.await?; } = code_receiver.await?;
oauth_state.handle_callback(&auth_code, &csrf_token).await?; oauth_state.handle_callback(&auth_code, &csrf_token).await?;
if let Err(e) = save_credentials(name, &oauth_state).await { let (client_id, token_response) = oauth_state.get_credentials().await?;
warn!("Failed to save credentials: {}", e);
}
let auth_manager = oauth_state let mut auth_manager = oauth_state
.into_authorization_manager() .into_authorization_manager()
.ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?; .ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?;
credential_store
.save(StoredCredentials {
client_id,
token_response,
})
.await?;
auth_manager.set_credential_store(credential_store);
Ok(auth_manager) Ok(auth_manager)
} }
+46 -63
View File
@@ -1,71 +1,54 @@
use oauth2::{basic::BasicTokenType, EmptyExtraTokenFields, StandardTokenResponse}; use rmcp::transport::auth::{AuthError, CredentialStore, StoredCredentials};
use reqwest::IntoUrl;
use rmcp::transport::{auth::OAuthState, AuthError};
use serde::{Deserialize, Serialize};
use crate::config::Config; use crate::config::Config;
#[derive(Debug, Clone, Serialize, Deserialize)] /// Goose-specific credential store that uses the Config system
pub struct SerializableCredentials { ///
pub client_id: String, /// This implementation stores OAuth credentials in the goose configuration
pub token_response: Option<StandardTokenResponse<EmptyExtraTokenFields, BasicTokenType>>, /// system, which handles secure storage (e.g., keychain integration).
#[derive(Clone)]
pub struct GooseCredentialStore {
name: String,
} }
fn secret_key(name: &str) -> String { impl GooseCredentialStore {
format!("oauth_creds_{name}") pub fn new(name: String) -> Self {
} Self { name }
}
pub async fn save_credentials( fn secret_key(&self) -> String {
name: &str, format!("oauth_creds_{}", self.name)
oauth_state: &OAuthState, }
) -> Result<(), Box<dyn std::error::Error>> { }
let config = Config::global();
let (client_id, token_response) = oauth_state.get_credentials().await?; #[async_trait::async_trait]
impl CredentialStore for GooseCredentialStore {
let credentials = SerializableCredentials { async fn load(&self) -> Result<Option<StoredCredentials>, AuthError> {
client_id, let config = Config::global();
token_response, let key = self.secret_key();
};
match config.get_secret::<StoredCredentials>(&key) {
let key = secret_key(name); Ok(credentials) => Ok(Some(credentials)),
config.set_secret(&key, &credentials)?; Err(_) => Ok(None), // No credentials found
}
Ok(()) }
}
async fn save(&self, credentials: StoredCredentials) -> Result<(), AuthError> {
async fn load_credentials( let config = Config::global();
name: &str, let key = self.secret_key();
) -> Result<SerializableCredentials, Box<dyn std::error::Error>> {
let config = Config::global(); config
let key = secret_key(name); .set_secret(&key, &credentials)
let credentials: SerializableCredentials = config.get_secret(&key)?; .map_err(|e| AuthError::InternalError(format!("Failed to save credentials: {}", e)))
}
Ok(credentials)
} async fn clear(&self) -> Result<(), AuthError> {
let config = Config::global();
pub fn clear_credentials(name: &str) -> Result<(), Box<dyn std::error::Error>> { let key = self.secret_key();
let config = Config::global();
config
Ok(config.delete_secret(&secret_key(name))?) .delete_secret(&key)
} .map_err(|e| AuthError::InternalError(format!("Failed to clear credentials: {}", e)))
pub async fn load_cached_state<U: IntoUrl>(
base_url: U,
name: &str,
) -> Result<OAuthState, AuthError> {
let credentials = load_credentials(name)
.await
.map_err(|e| AuthError::InternalError(format!("Failed to load credentials: {}", e)))?;
if let Some(token_response) = credentials.token_response {
let mut oauth_state = OAuthState::new(base_url, None).await?;
oauth_state
.set_credentials(&credentials.client_id, token_response)
.await?;
Ok(oauth_state)
} else {
Err(AuthError::InternalError(
"No token response in cached credentials".to_string(),
))
} }
} }
+4
View File
@@ -4643,6 +4643,10 @@
"inputSchema" "inputSchema"
], ],
"properties": { "properties": {
"_meta": {
"type": "object",
"additionalProperties": true
},
"annotations": { "annotations": {
"anyOf": [ "anyOf": [
{ {
+3
View File
@@ -810,6 +810,9 @@ export type TokenState = {
}; };
export type Tool = { export type Tool = {
_meta?: {
[key: string]: unknown;
};
annotations?: ToolAnnotations | { annotations?: ToolAnnotations | {
[key: string]: unknown; [key: string]: unknown;
}; };