diff --git a/Cargo.lock b/Cargo.lock index a5a2d38b4bde..58f8da859579 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.9.0", "schemars", "serde", "serde_json", @@ -2711,7 +2711,7 @@ dependencies = [ "once_cell", "paste", "regex", - "rmcp", + "rmcp 0.9.0", "serde", "serde_json", "tokio", @@ -2750,7 +2750,7 @@ dependencies = [ "open", "rand 0.8.5", "regex", - "rmcp", + "rmcp 0.9.0", "rustyline", "serde", "serde_json", @@ -2805,7 +2805,7 @@ dependencies = [ "rayon", "regex", "reqwest 0.11.27", - "rmcp", + "rmcp 0.8.5", "schemars", "serde", "serde_json", @@ -2856,7 +2856,7 @@ dependencies = [ "goose-mcp", "http 1.2.0", "reqwest 0.12.12", - "rmcp", + "rmcp 0.9.0", "schemars", "serde", "serde_json", @@ -5574,6 +5574,29 @@ 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", + "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", "chrono", "futures", @@ -5583,7 +5606,7 @@ dependencies = [ "pin-project-lite", "process-wrap", "reqwest 0.12.12", - "rmcp-macros", + "rmcp-macros 0.9.0", "schemars", "serde", "serde_json", @@ -5609,6 +5632,19 @@ dependencies = [ "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]] name = "ron" version = "0.8.1" diff --git a/Cargo.toml b/Cargo.toml index dc2bb1febc9a..4b6b8886460b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,7 +15,7 @@ uninlined_format_args = "allow" string_slice = "warn" [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.crates-io] diff --git a/crates/goose-bench/Cargo.toml b/crates/goose-bench/Cargo.toml index d91d09df32da..cabba6121a22 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 7fc20c75a57e..7404a2f84acc 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 34c0531c5e75..15ce91bc61af 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-server/Cargo.toml b/crates/goose-server/Cargo.toml index 59d4f2041823..98c43f542b77 100644 --- a/crates/goose-server/Cargo.toml +++ b/crates/goose-server/Cargo.toml @@ -50,5 +50,5 @@ path = "src/bin/generate_schema.rs" [dev-dependencies] tower = "0.5" -async-trait = "0.1" +async-trait = "0.1.89" tempfile = "3.15.0" diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index efc80530e8a3..b56c1ee61719 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 0401746ec7e4..dd8f356ded2f 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 d57fce564a27..1eccaa7223b9 100644 --- a/crates/goose/src/oauth/mod.rs +++ b/crates/goose/src/oauth/mod.rs @@ -1,9 +1,11 @@ +mod persist; + use axum::extract::{Query, State}; use axum::response::Html; use axum::routing::get; use axum::Router; use minijinja::render; -use rmcp::transport::auth::OAuthState; +use rmcp::transport::auth::{CredentialStore, OAuthState, StoredCredentials}; use rmcp::transport::AuthorizationManager; use serde::Deserialize; use std::net::SocketAddr; @@ -11,9 +13,7 @@ use std::sync::Arc; use tokio::sync::{oneshot, Mutex}; use tracing::warn; -use crate::oauth::persist::{clear_credentials, load_cached_state, save_credentials}; - -mod persist; +use crate::oauth::persist::GooseCredentialStore; const CALLBACK_TEMPLATE: &str = include_str!("oauth_callback.html"); @@ -32,18 +32,21 @@ 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 credential_store = GooseCredentialStore::new(name.clone()); + let mut auth_manager = AuthorizationManager::new(mcp_server_url).await?; + auth_manager.set_credential_store(credential_store.clone()); + + 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); } } + // No existing credentials or they were invalid - need to do the full oauth flow let (code_sender, code_receiver) = oneshot::channel::(); let app_state = AppState { 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 redirect_uri = format!("http://localhost:{}/oauth_callback", used_addr.port()); oauth_state .start_authorization(&[], redirect_uri.as_str(), Some("goose")) @@ -91,13 +95,20 @@ 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); - } + let (client_id, token_response) = oauth_state.get_credentials().await?; - let auth_manager = oauth_state + let mut auth_manager = oauth_state .into_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) } diff --git a/crates/goose/src/oauth/persist.rs b/crates/goose/src/oauth/persist.rs index 1600b8424862..b0c4155ecd02 100644 --- a/crates/goose/src/oauth/persist.rs +++ b/crates/goose/src/oauth/persist.rs @@ -1,71 +1,54 @@ -use oauth2::{basic::BasicTokenType, EmptyExtraTokenFields, StandardTokenResponse}; -use reqwest::IntoUrl; -use rmcp::transport::{auth::OAuthState, AuthError}; -use serde::{Deserialize, Serialize}; +use rmcp::transport::auth::{AuthError, CredentialStore, StoredCredentials}; use crate::config::Config; -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct SerializableCredentials { - pub client_id: String, - pub token_response: Option>, -} +/// Goose-specific credential store that uses the Config system +/// +/// This implementation stores OAuth credentials in the goose configuration +/// system, which handles secure storage (e.g., keychain integration). -fn secret_key(name: &str) -> String { - format!("oauth_creds_{name}") +#[derive(Clone)] +pub struct GooseCredentialStore { + 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?; - - let credentials = SerializableCredentials { - client_id, - token_response, - }; - - let key = secret_key(name); - config.set_secret(&key, &credentials)?; +impl GooseCredentialStore { + pub fn new(name: String) -> Self { + Self { name } + } - Ok(()) + fn secret_key(&self) -> String { + format!("oauth_creds_{}", self.name) + } } -async fn load_credentials( - name: &str, -) -> Result> { - let config = Config::global(); - let key = secret_key(name); - let credentials: SerializableCredentials = config.get_secret(&key)?; +#[async_trait::async_trait] +impl CredentialStore for GooseCredentialStore { + async fn load(&self) -> Result, AuthError> { + let config = Config::global(); + let key = self.secret_key(); - Ok(credentials) -} + match config.get_secret::(&key) { + Ok(credentials) => Ok(Some(credentials)), + Err(_) => Ok(None), // No credentials found + } + } -pub fn clear_credentials(name: &str) -> Result<(), Box> { - let config = Config::global(); + async fn save(&self, credentials: StoredCredentials) -> Result<(), AuthError> { + let config = Config::global(); + let key = self.secret_key(); - Ok(config.delete_secret(&secret_key(name))?) -} + config + .set_secret(&key, &credentials) + .map_err(|e| AuthError::InternalError(format!("Failed to save credentials: {}", e))) + } -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)))?; + async fn clear(&self) -> Result<(), AuthError> { + let config = Config::global(); + let key = self.secret_key(); - 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(), - )) + config + .delete_secret(&key) + .map_err(|e| AuthError::InternalError(format!("Failed to clear credentials: {}", e))) } } diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index e7e4f8a4f6da..8b87f06badb4 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 75de116e5690..c1e9b62014ac 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; };