ai-space/goose
3
0
Fork 0
mirror of https://github.com/aaif-goose/goose.git synced 2026-07-03 14:10:03 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix: implement a TokenStore which persists new refresh tokens

Browse Source
This commit is contained in:
Alex Hancock
2025-11-13 11:52:33 -05:00
parent 075cbb2a20
commit 93f930e86f
11 changed files with 131 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+42 -8
View File
@@ -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"
Cargo.toml
+2 -1
View File
@@ -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]
crates/goose-bench/Cargo.toml
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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/Cargo.toml
+1 -1
View File
@@ -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
+1
View File
@@ -674,6 +674,7 @@ impl ExtensionManager {
output_schema: tool.output_schema,
icons: None,
title: None,
meta: None,
});
}
}
crates/goose/src/oauth/mod.rs
+17 -12
View File
@@ -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<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 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::<CallbackParams>();
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"))?;
crates/goose/src/oauth/persist.rs
+58 -58
View File
@@ -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<StandardTokenResponse<EmptyExtraTokenFields, BasicTokenType>>,
struct SerializableCredentials {
token_response: StandardTokenResponse<EmptyExtraTokenFields, BasicTokenType>,
}
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<dyn std::error::Error>> {
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<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_token_{}", self.name)
}
}
#[async_trait]
impl TokenStore for ConfigTokenStore {
async fn load(
&self,
) -> Result<Option<OAuthTokenResponse>, Box<dyn std::error::Error + Send + Sync>> {
let config = Config::global();
let key = self.secret_key();
match config.get_secret::<SerializableCredentials>(&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<dyn std::error::Error + Send + Sync>> {
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<dyn std::error::Error + Send + Sync>> {
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<dyn TokenStore> {
Arc::new(ConfigTokenStore::new(name.to_string()))
}
ui/desktop/openapi.json
+4
View File
@@ -4643,6 +4643,10 @@
"inputSchema"
],
"properties": {
"_meta": {
"type": "object",
"additionalProperties": true
},
"annotations": {
"anyOf": [
{
ui/desktop/src/api/types.gen.ts
+3
View File
@@ -810,6 +810,9 @@ export type TokenState = {
};
export type Tool = {
_meta?: {
[key: string]: unknown;
};
annotations?: ToolAnnotations | {
[key: string]: unknown;
};