Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 44 additions & 8 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
2 changes: 1 addition & 1 deletion crates/goose-bench/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"] }
Expand Down
2 changes: 1 addition & 1 deletion crates/goose-cli/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"] }
Expand Down
2 changes: 1 addition & 1 deletion crates/goose-mcp/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
2 changes: 1 addition & 1 deletion crates/goose-server/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
2 changes: 1 addition & 1 deletion crates/goose/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
1 change: 1 addition & 0 deletions crates/goose/src/agents/extension_manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -674,6 +674,7 @@ impl ExtensionManager {
output_schema: tool.output_schema,
icons: None,
title: None,
meta: None,
});
}
}
Expand Down
39 changes: 25 additions & 14 deletions crates/goose/src/oauth/mod.rs
Original file line number Diff line number Diff line change
@@ -1,19 +1,19 @@
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;
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");

Expand All @@ -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))),
Expand Down Expand Up @@ -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"))
Expand All @@ -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)
}
93 changes: 38 additions & 55 deletions crates/goose/src/oauth/persist.rs
Original file line number Diff line number Diff line change
@@ -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).

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<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)?;
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<SerializableCredentials, Box<dyn std::error::Error>> {
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<Option<StoredCredentials>, AuthError> {
let config = Config::global();
let key = self.secret_key();

Ok(credentials)
}
match config.get_secret::<StoredCredentials>(&key) {
Ok(credentials) => Ok(Some(credentials)),
Err(_) => Ok(None), // No credentials found
}
}

pub fn clear_credentials(name: &str) -> Result<(), Box<dyn std::error::Error>> {
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<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)))?;
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)))
}
}
4 changes: 4 additions & 0 deletions ui/desktop/openapi.json
Original file line number Diff line number Diff line change
Expand Up @@ -4643,6 +4643,10 @@
"inputSchema"
],
"properties": {
"_meta": {
"type": "object",
"additionalProperties": true
},
"annotations": {
"anyOf": [
{
Expand Down
3 changes: 3 additions & 0 deletions ui/desktop/src/api/types.gen.ts
Original file line number Diff line number Diff line change
Expand Up @@ -810,6 +810,9 @@ export type TokenState = {
};

export type Tool = {
_meta?: {
[key: string]: unknown;
};
annotations?: ToolAnnotations | {
[key: string]: unknown;
};
Expand Down
Loading