fix: Implement a CredentialStore for auth (#5741)

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"

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]

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"] }

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"] }

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"

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"

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"

crates/goose/src/agents/extension_manager.rs

@@ -674,6 +674,7 @@ impl ExtensionManager {
                                 output_schema: tool.output_schema,
                                 icons: None,
                                 title: None,
+                                meta: None,
                             });
                         }
                     }

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<AuthorizationManager, anyhow::Error> {
-    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::<CallbackParams>();
     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)
 }

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<StandardTokenResponse<EmptyExtraTokenFields, BasicTokenType>>,
+/// 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).
+
+#[derive(Clone)]
+pub struct GooseCredentialStore {
+    name: String,
 }
 
-fn secret_key(name: &str) -> String {
-    format!("oauth_creds_{name}")
-}
+impl GooseCredentialStore {
+    pub fn new(name: String) -> Self {
+        Self { name }
+    }
 
-pub async fn save_credentials(
-    name: &str,
-    oauth_state: &OAuthState,
-) -> Result<(), Box<dyn std::error::Error>> {
-    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)?;
-
-    Ok(())
-}
-
-async fn load_credentials(
-    name: &str,
-) -> Result<SerializableCredentials, Box<dyn std::error::Error>> {
-    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<dyn std::error::Error>> {
-    let config = Config::global();
-
-    Ok(config.delete_secret(&secret_key(name))?)
-}
-
-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(),
-        ))
+    fn secret_key(&self) -> String {
+        format!("oauth_creds_{}", self.name)
+    }
+}
+
+#[async_trait::async_trait]
+impl CredentialStore for GooseCredentialStore {
+    async fn load(&self) -> Result<Option<StoredCredentials>, AuthError> {
+        let config = Config::global();
+        let key = self.secret_key();
+
+        match config.get_secret::<StoredCredentials>(&key) {
+            Ok(credentials) => Ok(Some(credentials)),
+            Err(_) => Ok(None), // No credentials found
+        }
+    }
+
+    async fn save(&self, credentials: StoredCredentials) -> Result<(), AuthError> {
+        let config = Config::global();
+        let key = self.secret_key();
+
+        config
+            .set_secret(&key, &credentials)
+            .map_err(|e| AuthError::InternalError(format!("Failed to save credentials: {}", e)))
+    }
+
+    async fn clear(&self) -> Result<(), AuthError> {
+        let config = Config::global();
+        let key = self.secret_key();
+
+        config
+            .delete_secret(&key)
+            .map_err(|e| AuthError::InternalError(format!("Failed to clear credentials: {}", e)))
     }
 }

ui/desktop/openapi.json

@@ -4643,6 +4643,10 @@
           "inputSchema"
         ],
         "properties": {
+          "_meta": {
+            "type": "object",
+            "additionalProperties": true
+          },
           "annotations": {
             "anyOf": [
               {

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;
     };