diff --git a/Cargo.lock b/Cargo.lock index d969263452..1b6fe6c7ac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -281,9 +281,9 @@ dependencies = [ [[package]] name = "async-trait" -version = "0.1.87" +version = "0.1.89" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d556ec1359574147ec0c4fc5eb525f3f23263a592b1a9c07e0a75b427de55c97" +checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" dependencies = [ "proc-macro2", "quote", @@ -2661,7 +2661,7 @@ dependencies = [ "rand 0.8.5", "regex", "reqwest 0.12.12", - "rmcp", + "rmcp 0.8.5 (git+https://github.com/modelcontextprotocol/rust-sdk?branch=alexhancock%2Ftoken-store)", "schemars", "serde", "serde_json", @@ -2711,7 +2711,7 @@ dependencies = [ "once_cell", "paste", "regex", - "rmcp", + "rmcp 0.8.5 (git+https://github.com/modelcontextprotocol/rust-sdk?branch=alexhancock%2Ftoken-store)", "serde", "serde_json", "tokio", @@ -2750,7 +2750,7 @@ dependencies = [ "open", "rand 0.8.5", "regex", - "rmcp", + "rmcp 0.8.5 (git+https://github.com/modelcontextprotocol/rust-sdk?branch=alexhancock%2Ftoken-store)", "rustyline", "serde", "serde_json", @@ -2804,7 +2804,7 @@ dependencies = [ "rayon", "regex", "reqwest 0.11.27", - "rmcp", + "rmcp 0.8.5 (registry+https://github.com/rust-lang/crates.io-index)", "schemars", "serde", "serde_json", @@ -2855,7 +2855,7 @@ dependencies = [ "goose-mcp", "http 1.2.0", "reqwest 0.12.12", - "rmcp", + "rmcp 0.8.5 (git+https://github.com/modelcontextprotocol/rust-sdk?branch=alexhancock%2Ftoken-store)", "schemars", "serde", "serde_json", @@ -5573,6 +5573,28 @@ version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5947688160b56fb6c827e3c20a72c90392a1d7e9dec74749197aa1780ac42ca" dependencies = [ + "base64 0.22.1", + "chrono", + "futures", + "paste", + "pin-project-lite", + "rmcp-macros 0.8.5 (registry+https://github.com/rust-lang/crates.io-index)", + "schemars", + "serde", + "serde_json", + "thiserror 2.0.12", + "tokio", + "tokio-stream", + "tokio-util", + "tracing", +] + +[[package]] +name = "rmcp" +version = "0.8.5" +source = "git+https://github.com/modelcontextprotocol/rust-sdk?branch=alexhancock%2Ftoken-store#3c1d7ee9fd8a37d497c8a9be9fc1e104d9adc1fd" +dependencies = [ + "async-trait", "base64 0.22.1", "chrono", "futures", @@ -5582,7 +5604,7 @@ dependencies = [ "pin-project-lite", "process-wrap", "reqwest 0.12.12", - "rmcp-macros", + "rmcp-macros 0.8.5 (git+https://github.com/modelcontextprotocol/rust-sdk?branch=alexhancock%2Ftoken-store)", "schemars", "serde", "serde_json", @@ -5608,6 +5630,18 @@ dependencies = [ "syn 2.0.99", ] +[[package]] +name = "rmcp-macros" +version = "0.8.5" +source = "git+https://github.com/modelcontextprotocol/rust-sdk?branch=alexhancock%2Ftoken-store#3c1d7ee9fd8a37d497c8a9be9fc1e104d9adc1fd" +dependencies = [ + "darling 0.21.0", + "proc-macro2", + "quote", + "serde_json", + "syn 2.0.99", +] + [[package]] name = "ron" version = "0.8.1" diff --git a/Cargo.toml b/Cargo.toml index dc2bb1febc..96da08750e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,7 +15,8 @@ uninlined_format_args = "allow" string_slice = "warn" [workspace.dependencies] -rmcp = { version = "0.8.5", features = ["schemars", "auth"] } +# rmcp = { version = "0.8.5", features = ["schemars", "auth"] } +rmcp = { git = "https://github.com/modelcontextprotocol/rust-sdk", branch = "alexhancock/token-store", features = ["schemars", "auth"] } # Patch for Windows cross-compilation issue with crunchy [patch.crates-io] diff --git a/crates/goose-bench/Cargo.toml b/crates/goose-bench/Cargo.toml index d91d09df32..cabba6121a 100644 --- a/crates/goose-bench/Cargo.toml +++ b/crates/goose-bench/Cargo.toml @@ -16,7 +16,7 @@ paste = "1.0" ctor = "0.2.7" goose = { path = "../goose" } rmcp = { workspace = true } -async-trait = "0.1.86" +async-trait = "0.1.89" chrono = { version = "0.4", features = ["serde"] } serde_json = "1.0" serde = { version = "1.0", features = ["derive"] } diff --git a/crates/goose-cli/Cargo.toml b/crates/goose-cli/Cargo.toml index dbb5468437..e50c5fce71 100644 --- a/crates/goose-cli/Cargo.toml +++ b/crates/goose-cli/Cargo.toml @@ -43,7 +43,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt", "json", tracing-appender = "0.2" once_cell = "1.20.2" shlex = "1.3.0" -async-trait = "0.1.86" +async-trait = "0.1.89" base64 = "0.22.1" regex = "1.11.1" nix = { version = "0.30.1", features = ["process", "signal"] } diff --git a/crates/goose-mcp/Cargo.toml b/crates/goose-mcp/Cargo.toml index 34c0531c5e..15ce91bc61 100644 --- a/crates/goose-mcp/Cargo.toml +++ b/crates/goose-mcp/Cargo.toml @@ -33,7 +33,7 @@ reqwest = { version = "0.11", features = [ "json", "rustls-tls-native-roots", ], default-features = false } -async-trait = "0.1" +async-trait = "0.1.89" chrono = { version = "0.4.38", features = ["serde"] } etcetera = "0.8.0" tempfile = "3.8" diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index efc80530e8..b56c1ee617 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -49,7 +49,7 @@ serde_urlencoded = "0.7" jsonschema = "0.30.0" uuid = { version = "1.0", features = ["v4"] } regex = "1.11.1" -async-trait = "0.1" +async-trait = "0.1.89" async-stream = "0.3" minijinja = { version = "2.10.2", features = ["loader"] } include_dir = "0.7.4" diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index 3ae559fcae..b8c0fe3a53 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -674,6 +674,7 @@ impl ExtensionManager { output_schema: tool.output_schema, icons: None, title: None, + meta: None, }); } } diff --git a/crates/goose/src/oauth/mod.rs b/crates/goose/src/oauth/mod.rs index d57fce564a..29999cbea2 100644 --- a/crates/goose/src/oauth/mod.rs +++ b/crates/goose/src/oauth/mod.rs @@ -11,7 +11,7 @@ use std::sync::Arc; use tokio::sync::{oneshot, Mutex}; use tracing::warn; -use crate::oauth::persist::{clear_credentials, load_cached_state, save_credentials}; +use crate::oauth::persist::create_token_store; mod persist; @@ -32,18 +32,24 @@ pub async fn oauth_flow( mcp_server_url: &String, name: &String, ) -> Result { - if let Ok(oauth_state) = load_cached_state(mcp_server_url, name).await { - if let Some(authorization_manager) = oauth_state.into_authorization_manager() { - if authorization_manager.refresh_token().await.is_ok() { - return Ok(authorization_manager); - } + let token_store = create_token_store(name); + + // Try to create an AuthorizationManager with the existing token store + // If we have a valid token, refresh_token will succeed + if let Ok(auth_manager) = + AuthorizationManager::with_token_store(mcp_server_url, token_store.clone()).await + { + if auth_manager.refresh_token().await.is_ok() { + return Ok(auth_manager); } - if let Err(e) = clear_credentials(name) { + // Token is invalid or expired, clear it + if let Err(e) = token_store.clear().await { warn!("error clearing bad credentials: {}", e); } } + // Need to do full OAuth flow let (code_sender, code_receiver) = oneshot::channel::(); let app_state = AppState { code_receiver: Arc::new(Mutex::new(Some(code_sender))), @@ -73,7 +79,9 @@ pub async fn oauth_flow( } }); - let mut oauth_state = OAuthState::new(mcp_server_url, None).await?; + // Use OAuthState with the token store for the authorization flow + let auth_manager = AuthorizationManager::with_token_store(mcp_server_url, token_store).await?; + let mut oauth_state = OAuthState::Unauthorized(auth_manager); let redirect_uri = format!("http://localhost:{}/oauth_callback", used_addr.port()); oauth_state .start_authorization(&[], redirect_uri.as_str(), Some("goose")) @@ -91,10 +99,7 @@ pub async fn oauth_flow( } = code_receiver.await?; oauth_state.handle_callback(&auth_code, &csrf_token).await?; - if let Err(e) = save_credentials(name, &oauth_state).await { - warn!("Failed to save credentials: {}", e); - } - + // The token is automatically saved via the TokenStore during handle_callback let auth_manager = oauth_state .into_authorization_manager() .ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?; diff --git a/crates/goose/src/oauth/persist.rs b/crates/goose/src/oauth/persist.rs index 1600b84248..a51b6a30aa 100644 --- a/crates/goose/src/oauth/persist.rs +++ b/crates/goose/src/oauth/persist.rs @@ -1,71 +1,71 @@ +use async_trait::async_trait; use oauth2::{basic::BasicTokenType, EmptyExtraTokenFields, StandardTokenResponse}; -use reqwest::IntoUrl; -use rmcp::transport::{auth::OAuthState, AuthError}; +use rmcp::transport::auth::{OAuthTokenResponse, TokenStore}; use serde::{Deserialize, Serialize}; +use std::sync::Arc; use crate::config::Config; +/// Credentials stored for an OAuth provider #[derive(Debug, Clone, Serialize, Deserialize)] -pub struct SerializableCredentials { - pub client_id: String, - pub token_response: Option>, +struct SerializableCredentials { + token_response: StandardTokenResponse, } -fn secret_key(name: &str) -> String { - format!("oauth_creds_{name}") +/// Token store implementation that uses goose's Config secret storage +pub struct ConfigTokenStore { + name: String, } -pub async fn save_credentials( - name: &str, - oauth_state: &OAuthState, -) -> Result<(), Box> { - let config = Config::global(); - let (client_id, token_response) = oauth_state.get_credentials().await?; +impl ConfigTokenStore { + pub fn new(name: String) -> Self { + Self { name } + } - let credentials = SerializableCredentials { - client_id, - token_response, - }; - - let key = secret_key(name); - config.set_secret(&key, &credentials)?; - - Ok(()) -} - -async fn load_credentials( - name: &str, -) -> Result> { - let config = Config::global(); - let key = secret_key(name); - let credentials: SerializableCredentials = config.get_secret(&key)?; - - Ok(credentials) -} - -pub fn clear_credentials(name: &str) -> Result<(), Box> { - let config = Config::global(); - - Ok(config.delete_secret(&secret_key(name))?) -} - -pub async fn load_cached_state( - base_url: U, - name: &str, -) -> Result { - 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(), - )) + fn secret_key(&self) -> String { + format!("oauth_token_{}", self.name) } } + +#[async_trait] +impl TokenStore for ConfigTokenStore { + async fn load( + &self, + ) -> Result, Box> { + let config = Config::global(); + let key = self.secret_key(); + + match config.get_secret::(&key) { + Ok(credentials) => Ok(Some(credentials.token_response)), + Err(crate::config::ConfigError::NotFound(_)) => Ok(None), + Err(e) => Err(Box::new(e)), + } + } + + async fn save( + &self, + token: &OAuthTokenResponse, + ) -> Result<(), Box> { + let config = Config::global(); + let key = self.secret_key(); + + let credentials = SerializableCredentials { + token_response: token.clone(), + }; + + config.set_secret(&key, &credentials)?; + Ok(()) + } + + async fn clear(&self) -> Result<(), Box> { + let config = Config::global(); + let key = self.secret_key(); + config.delete_secret(&key)?; + Ok(()) + } +} + +/// Create a TokenStore for use with AuthorizationManager +pub fn create_token_store(name: &str) -> Arc { + Arc::new(ConfigTokenStore::new(name.to_string())) +} diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index e7e4f8a4f6..8b87f06bad 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -4643,6 +4643,10 @@ "inputSchema" ], "properties": { + "_meta": { + "type": "object", + "additionalProperties": true + }, "annotations": { "anyOf": [ { diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index 75de116e56..c1e9b62014 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -810,6 +810,9 @@ export type TokenState = { }; export type Tool = { + _meta?: { + [key: string]: unknown; + }; annotations?: ToolAnnotations | { [key: string]: unknown; };