diff --git a/.gitignore b/.gitignore index 87207b0..ba01919 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,6 @@ .idea/ /coyote.iml /.idea/ -.coyote +.coyote/** +.sisyphus/** +.coyote-project.json diff --git a/src/cli/mod.rs b/src/cli/mod.rs index b57859b..9f86023 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -765,10 +765,7 @@ mod tests { assert_eq!(cli.mcp_add, Some("notion".to_string())); assert!(matches!(cli.transport, Some(McpTransportArg::Http))); assert_eq!(cli.url, Some("https://mcp.notion.com/mcp".to_string())); - assert_eq!( - cli.header, - vec!["Authorization: Bearer {{NOTION_TOKEN}}"] - ); + assert_eq!(cli.header, vec!["Authorization: Bearer {{NOTION_TOKEN}}"]); assert!(cli.mcp_command.is_empty()); } diff --git a/src/client/oauth.rs b/src/client/oauth.rs index d45fdd5..60d86ba 100644 --- a/src/client/oauth.rs +++ b/src/client/oauth.rs @@ -2,13 +2,13 @@ use super::access_token::{is_valid_access_token, set_access_token}; use super::openai_compatible_oauth::OpenAICompatibleOAuthProvider; use super::{ClientConfig, ProviderModels}; use crate::config::paths; -use anyhow::{Context, Result, anyhow, bail}; +use anyhow::{Context, Error, Result, anyhow, bail}; use base64::Engine; use base64::engine::general_purpose::URL_SAFE_NO_PAD; use chrono::Utc; use indexmap::IndexMap; use inquire::Text; -use reqwest::{Client as ReqwestClient, RequestBuilder}; +use reqwest::{Client as ReqwestClient, RequestBuilder, StatusCode}; use serde::{Deserialize, Serialize}; use serde_json::Value; use sha2::{Digest, Sha256}; @@ -16,6 +16,10 @@ use std::collections::HashMap; use std::fs; use std::io::{BufRead, BufReader, Write}; use std::net::TcpListener; +use std::path::PathBuf; +use std::sync::{Arc, OnceLock}; +use std::time::Duration; +use tokio::sync; use url::Url; use uuid::Uuid; @@ -197,10 +201,20 @@ pub struct OAuthTokens { pub account_id: Option, } +const TOKEN_ENDPOINT_TIMEOUT: Duration = Duration::from_secs(30); + pub async fn run_oauth_flow(provider: &dyn OAuthProvider, client_name: &str) -> Result<()> { match provider.flow() { OAuthFlow::Pkce => run_pkce_flow(provider, client_name).await, - OAuthFlow::ClientCredentials => run_client_credentials_flow(provider, client_name).await, + OAuthFlow::ClientCredentials => { + run_client_credentials_flow(provider, client_name).await?; + println!( + "Successfully authenticated client '{}' with {} via OAuth (client_credentials). Tokens saved.", + client_name, + provider.provider_name() + ); + Ok(()) + } OAuthFlow::DeviceCode => run_device_code_flow(provider, client_name).await, } } @@ -301,12 +315,20 @@ async fn run_pkce_flow(provider: &dyn OAuthProvider, client_name: &str) -> Resul let access_token = response["access_token"] .as_str() - .ok_or_else(|| anyhow!("Missing access_token in response: {response}"))? + .ok_or_else(|| { + anyhow!( + "Missing access_token in response (keys: {})", + token_response_keys(&response) + ) + })? .to_string(); let refresh_token = response["refresh_token"].as_str().map(|s| s.to_string()); - let expires_in = response["expires_in"] - .as_i64() - .ok_or_else(|| anyhow!("Missing expires_in in response: {response}"))?; + let expires_in = response["expires_in"].as_i64().ok_or_else(|| { + anyhow!( + "Missing expires_in in response (keys: {})", + token_response_keys(&response) + ) + })?; let expires_at = Utc::now().timestamp() + expires_in; @@ -334,7 +356,9 @@ async fn run_client_credentials_flow( provider: &dyn OAuthProvider, client_name: &str, ) -> Result<()> { - let client = ReqwestClient::new(); + let client = ReqwestClient::builder() + .timeout(TOKEN_ENDPOINT_TIMEOUT) + .build()?; let scopes = provider.scopes(); let mut params: Vec<(&str, &str)> = vec![ ("grant_type", "client_credentials"), @@ -349,11 +373,19 @@ async fn run_client_credentials_flow( let access_token = response["access_token"] .as_str() - .ok_or_else(|| anyhow!("Missing access_token in client_credentials response: {response}"))? + .ok_or_else(|| { + anyhow!( + "Missing access_token in client_credentials response (keys: {})", + token_response_keys(&response) + ) + })? .to_string(); - let expires_in = response["expires_in"] - .as_i64() - .ok_or_else(|| anyhow!("Missing expires_in in client_credentials response: {response}"))?; + let expires_in = response["expires_in"].as_i64().ok_or_else(|| { + anyhow!( + "Missing expires_in in client_credentials response (keys: {})", + token_response_keys(&response) + ) + })?; let expires_at = Utc::now().timestamp() + expires_in; let tokens = OAuthTokens { @@ -363,11 +395,6 @@ async fn run_client_credentials_flow( account_id: provider.extract_account_id(&response), }; save_oauth_tokens(client_name, &tokens)?; - println!( - "Successfully authenticated client '{}' with {} via OAuth (client_credentials). Tokens saved.", - client_name, - provider.provider_name() - ); Ok(()) } @@ -417,19 +444,28 @@ async fn run_device_code_flow(provider: &dyn OAuthProvider, client_name: &str) - let device_code = device_response["device_code"] .as_str() .ok_or_else(|| { - anyhow!("Missing device_code in device authorization response: {device_response}") + anyhow!( + "Missing device_code in device authorization response (keys: {})", + token_response_keys(&device_response) + ) })? .to_string(); let user_code = device_response["user_code"] .as_str() .ok_or_else(|| { - anyhow!("Missing user_code in device authorization response: {device_response}") + anyhow!( + "Missing user_code in device authorization response (keys: {})", + token_response_keys(&device_response) + ) })? .to_string(); let verification_uri = device_response["verification_uri"] .as_str() .ok_or_else(|| { - anyhow!("Missing verification_uri in device authorization response: {device_response}") + anyhow!( + "Missing verification_uri in device authorization response (keys: {})", + token_response_keys(&device_response) + ) })? .to_string(); let verification_uri_complete = device_response["verification_uri_complete"] @@ -551,10 +587,76 @@ fn save_oauth_tokens(client_name: &str, tokens: &OAuthTokens) -> Result<()> { fs::create_dir_all(parent)?; } let json = serde_json::to_string_pretty(tokens)?; - fs::write(path, json)?; + // Write-then-rename so a crash mid-write never truncates the live token file. + let mut tmp = path.clone().into_os_string(); + tmp.push(".tmp"); + let tmp = PathBuf::from(tmp); + // Tokens are live credentials: create the file owner-only, not umask-default. + let mut options = fs::OpenOptions::new(); + options.write(true).create(true).truncate(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.mode(0o600); + } + options.open(&tmp)?.write_all(json.as_bytes())?; + fs::rename(&tmp, &path)?; Ok(()) } +pub(crate) fn token_response_keys(response: &Value) -> String { + match response.as_object() { + Some(map) => { + let keys: Vec<&str> = map.keys().map(String::as_str).collect(); + format!("[{}]", keys.join(", ")) + } + None => "".to_string(), + } +} + +fn parse_refresh_response( + status: StatusCode, + response: &Value, + previous_refresh_token: Option<&str>, +) -> Result<(String, Option, i64)> { + if let Some(error) = response["error"].as_str() { + let description = response["error_description"] + .as_str() + .unwrap_or("no description"); + if matches!(error, "invalid_grant" | "invalid_token") { + bail!( + "OAuth refresh token was rejected ({error}: {description}). Please re-authenticate." + ); + } + bail!("Token refresh failed ({error}: {description})"); + } + if !status.is_success() { + bail!("Token refresh failed with HTTP status {status}"); + } + + let access_token = response["access_token"] + .as_str() + .ok_or_else(|| { + anyhow!( + "Missing access_token in refresh response (keys: {})", + token_response_keys(response) + ) + })? + .to_string(); + let refresh_token = response["refresh_token"] + .as_str() + .or(previous_refresh_token) + .map(str::to_string); + let expires_in = response["expires_in"].as_i64().ok_or_else(|| { + anyhow!( + "Missing expires_in in refresh response (keys: {})", + token_response_keys(response) + ) + })?; + + Ok((access_token, refresh_token, expires_in)) +} + pub async fn refresh_oauth_token( client: &ReqwestClient, provider: &dyn OAuthProvider, @@ -577,19 +679,23 @@ pub async fn refresh_oauth_token( ], ); - let response: Value = request.send().await?.json().await?; + let (status, response) = tokio::time::timeout(TOKEN_ENDPOINT_TIMEOUT, async { + let response = request.send().await?; + let status = response.status(); + let body: Value = response.json().await?; + Ok::<_, Error>((status, body)) + }) + .await + .map_err(|_| { + anyhow!( + "Token refresh for '{}' timed out after {}s", + client_name, + TOKEN_ENDPOINT_TIMEOUT.as_secs() + ) + })??; - let access_token = response["access_token"] - .as_str() - .ok_or_else(|| anyhow!("Missing access_token in refresh response: {response}"))? - .to_string(); - let refresh_token = response["refresh_token"] - .as_str() - .map(|s| s.to_string()) - .or_else(|| tokens.refresh_token.clone()); - let expires_in = response["expires_in"] - .as_i64() - .ok_or_else(|| anyhow!("Missing expires_in in refresh response: {response}"))?; + let (access_token, refresh_token, expires_in) = + parse_refresh_response(status, &response, tokens.refresh_token.as_deref())?; let expires_at = Utc::now().timestamp() + expires_in; @@ -609,6 +715,20 @@ pub async fn refresh_oauth_token( Ok(new_tokens) } +/// Per-client lock so concurrent requests perform a single refresh. +/// Returns a clone of the Arc so the parking_lot guard is dropped before the +/// caller awaits on the tokio mutex. +fn refresh_guard(client_name: &str) -> Arc> { + static GUARDS: OnceLock>>>> = + OnceLock::new(); + GUARDS + .get_or_init(Default::default) + .lock() + .entry(client_name.to_string()) + .or_default() + .clone() +} + pub async fn prepare_oauth_access_token( client: &ReqwestClient, provider: &dyn OAuthProvider, @@ -624,15 +744,35 @@ pub async fn prepare_oauth_access_token( }; let tokens = if Utc::now().timestamp() >= tokens.expires_at { - match provider.flow() { - OAuthFlow::Pkce | OAuthFlow::DeviceCode => { - refresh_oauth_token(client, provider, client_name, &tokens).await? - } - OAuthFlow::ClientCredentials => { - run_client_credentials_flow(provider, client_name).await?; - load_oauth_tokens(client_name) - .ok_or_else(|| anyhow!("Token file missing after client_credentials refresh"))? + let guard = refresh_guard(client_name); + let _guard = guard.lock().await; + + // A concurrent caller may have refreshed while we waited for the + // lock; a valid in-memory token means the winner already populated + // the cache. + if is_valid_access_token(client_name) { + return Ok(true); + } + + let tokens = match load_oauth_tokens(client_name) { + Some(t) => t, + None => return Ok(false), + }; + + if Utc::now().timestamp() >= tokens.expires_at { + match provider.flow() { + OAuthFlow::Pkce | OAuthFlow::DeviceCode => { + refresh_oauth_token(client, provider, client_name, &tokens).await? + } + OAuthFlow::ClientCredentials => { + run_client_credentials_flow(provider, client_name).await?; + load_oauth_tokens(client_name).ok_or_else(|| { + anyhow!("Token file missing after client_credentials refresh") + })? + } } + } else { + tokens } } else { tokens @@ -886,11 +1026,54 @@ pub(crate) fn client_config_info( #[cfg(test)] mod tests { + use std::ffi::OsString; + use std::path::PathBuf; use std::str; + use std::time::UNIX_EPOCH; use super::*; use crate::client::openai_compatible::OpenAICompatibleConfig; use crate::client::{ModelData, ProviderModels}; + use crate::utils::get_env_name; + use serial_test::serial; + use std::{env, time::SystemTime}; + + fn with_temp_cache(f: F) { + struct Restore { + key: String, + prev: Option, + root: PathBuf, + } + impl Drop for Restore { + fn drop(&mut self) { + unsafe { + match self.prev.take() { + Some(v) => env::set_var(&self.key, v), + None => env::remove_var(&self.key), + } + } + let _ = fs::remove_dir_all(&self.root); + } + } + + let unique = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + let root = env::temp_dir().join(format!("coyote-client-oauth-test-{unique}")); + fs::create_dir_all(&root).unwrap(); + let env_key = get_env_name("cache_dir"); + let prev = env::var_os(&env_key); + unsafe { + env::set_var(&env_key, &root); + } + let _restore = Restore { + key: env_key, + prev, + root, + }; + f(); + } fn base_config() -> OAuthConfig { OAuthConfig { @@ -1468,4 +1651,104 @@ scopes: "body missing grant_type param: {body}" ); } + + #[test] + #[serial] + fn save_oauth_tokens_roundtrips_and_leaves_no_tmp_file() { + with_temp_cache(|| { + let tokens = OAuthTokens { + access_token: "at-123".into(), + refresh_token: Some("rt-456".into()), + expires_at: 1234567890, + account_id: Some("acct-789".into()), + }; + + save_oauth_tokens("atomic-test", &tokens).unwrap(); + + let loaded = load_oauth_tokens("atomic-test").unwrap(); + assert_eq!(loaded.access_token, "at-123"); + assert_eq!(loaded.refresh_token.as_deref(), Some("rt-456")); + assert_eq!(loaded.expires_at, 1234567890); + assert_eq!(loaded.account_id.as_deref(), Some("acct-789")); + + let dir = paths::oauth_tokens_dir(); + let leftover_tmp = fs::read_dir(&dir) + .unwrap() + .any(|e| e.unwrap().file_name().to_string_lossy().ends_with(".tmp")); + assert!(!leftover_tmp, "temp file left behind in {dir:?}"); + + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let mode = fs::metadata(paths::token_file("atomic-test")) + .unwrap() + .permissions() + .mode(); + assert_eq!(mode & 0o777, 0o600, "token file mode was {mode:o}"); + } + }); + } + + #[test] + fn parse_refresh_response_invalid_grant_redacts_and_prompts_reauth() { + let response = serde_json::json!({ + "error": "invalid_grant", + "error_description": "refresh token revoked", + "refresh_token": "planted-secret-token", + }); + + let err = parse_refresh_response(StatusCode::BAD_REQUEST, &response, Some("old-rt")) + .unwrap_err() + .to_string(); + + assert!(err.contains("re-authenticate"), "unexpected error: {err}"); + assert!( + !err.contains("planted-secret-token"), + "error leaked token material: {err}" + ); + } + + #[test] + fn token_response_keys_lists_keys_without_values() { + let response = serde_json::json!({ + "access_token": "secret-at", + "token_type": "SecretBearer", + }); + + let keys = token_response_keys(&response); + + assert!(keys.contains("access_token"), "missing key name: {keys}"); + assert!(keys.contains("token_type"), "missing key name: {keys}"); + assert!(!keys.contains("secret-at"), "leaked value: {keys}"); + assert!(!keys.contains("SecretBearer"), "leaked value: {keys}"); + } + + #[test] + fn parse_refresh_response_rotates_refresh_token_when_present() { + let response = serde_json::json!({ + "access_token": "new-at", + "refresh_token": "new-rt", + "expires_in": 3600, + }); + + let (access_token, refresh_token, expires_in) = + parse_refresh_response(StatusCode::OK, &response, Some("old-rt")).unwrap(); + + assert_eq!(access_token, "new-at"); + assert_eq!(refresh_token.as_deref(), Some("new-rt")); + assert_eq!(expires_in, 3600); + } + + #[test] + fn parse_refresh_response_keeps_old_refresh_token_when_absent() { + let response = serde_json::json!({ + "access_token": "new-at", + "expires_in": 3600, + }); + + let (_, refresh_token, _) = + parse_refresh_response(StatusCode::OK, &response, Some("old-rt")).unwrap(); + + assert_eq!(refresh_token.as_deref(), Some("old-rt")); + } } diff --git a/src/config/mcp_factory.rs b/src/config/mcp_factory.rs index 872cfd3..ef9f3f4 100644 --- a/src/config/mcp_factory.rs +++ b/src/config/mcp_factory.rs @@ -103,7 +103,7 @@ impl McpFactory { } let bearer_token = if spec.is_remote() { - oauth::load_valid_mcp_token(name) + oauth::load_or_refresh_mcp_token(name).await } else { None }; diff --git a/src/mcp/mod.rs b/src/mcp/mod.rs index 9cd7bad..b826b85 100644 --- a/src/mcp/mod.rs +++ b/src/mcp/mod.rs @@ -327,7 +327,7 @@ impl McpRegistry { .with_context(|| format!("MCP server not found in config: {id}"))?; let bearer_token = if spec.is_remote() { - oauth::load_valid_mcp_token(&id) + oauth::load_or_refresh_mcp_token(&id).await } else { None }; diff --git a/src/mcp/oauth.rs b/src/mcp/oauth.rs index 3455251..99b15f4 100644 --- a/src/mcp/oauth.rs +++ b/src/mcp/oauth.rs @@ -1,15 +1,25 @@ -use crate::client::oauth::{OAuthProvider, TokenRequestFormat, load_oauth_tokens, run_oauth_flow}; +use crate::client::oauth::{ + OAuthProvider, OAuthTokens, TokenRequestFormat, load_oauth_tokens, refresh_oauth_token, + run_oauth_flow, token_response_keys, +}; use crate::config::paths; use anyhow::{Context, Result, anyhow}; use chrono::Utc; use inquire::Text; -use log::warn; +use log::{debug, warn}; use reqwest::Client; use serde::{Deserialize, Serialize}; +use std::collections::HashMap; use std::fs; use std::net::TcpListener; +use std::sync::{Arc, OnceLock}; +use std::time::{Duration, Instant}; +use tokio::sync; use url::Url; +const REFRESH_HTTP_TIMEOUT: Duration = Duration::from_secs(10); +const REFRESH_FAILURE_BACKOFF: Duration = Duration::from_secs(60); + #[derive(Debug, Deserialize)] struct ProtectedResourceMetadata { #[serde(default)] @@ -34,6 +44,10 @@ struct McpRegistration { client_id: String, #[serde(default)] redirect_uri: Option, + #[serde(default)] + token_url: Option, + #[serde(default)] + resource: Option, } struct DiscoveredOAuth { @@ -124,8 +138,19 @@ pub async fn run_mcp_oauth_flow( None }; - let (client_id, redirect_uri) = if let Some(reused) = cached_reuse { - reused + let (client_id, redirect_uri) = if let Some((client_id, redirect_uri)) = cached_reuse { + // Re-save so registrations cached before token_url/resource were + // persisted gain them, enabling token refresh next time. + if let Err(e) = save_registration( + server_name, + &client_id, + &redirect_uri, + &metadata.token_endpoint, + &resource, + ) { + debug!("Failed to update cached MCP registration for '{server_name}': {e}"); + } + (client_id, redirect_uri) } else { let bind_addr = format!("127.0.0.1:{}", callback_port.unwrap_or(0)); let listener = TcpListener::bind(&bind_addr)?; @@ -137,10 +162,7 @@ pub async fn run_mcp_oauth_flow( id.to_string() } else if let Some(reg_endpoint) = &metadata.registration_endpoint { match register_client(reg_endpoint, &redirect_uri).await { - Ok(id) => { - let _ = save_registration(server_name, &id, &redirect_uri); - id - } + Ok(id) => id, Err(e) => { warn!("Dynamic client registration failed: {e}. Falling back to manual entry."); Text::new("Enter the OAuth client ID for this MCP server:") @@ -153,6 +175,18 @@ pub async fn run_mcp_oauth_flow( .prompt() .context("Failed to read client ID")? }; + // Persist regardless of how the client_id was obtained (DCR, config, + // or manual entry) so refresh_mcp_token can run the refresh_token + // grant later without interactive re-auth. + if let Err(e) = save_registration( + server_name, + &client_id, + &redirect_uri, + &metadata.token_endpoint, + &resource, + ) { + debug!("Failed to cache MCP registration for '{server_name}': {e}"); + } (client_id, redirect_uri) }; @@ -168,12 +202,113 @@ pub async fn run_mcp_oauth_flow( run_oauth_flow(&provider, &mcp_token_key(server_name)).await } -pub fn load_valid_mcp_token(server_name: &str) -> Option { - let tokens = load_oauth_tokens(&mcp_token_key(server_name))?; +pub async fn load_or_refresh_mcp_token(server_name: &str) -> Option { + let key = mcp_token_key(server_name); + let tokens = load_oauth_tokens(&key)?; if Utc::now().timestamp() < tokens.expires_at { - Some(tokens.access_token) - } else { - None + return Some(tokens.access_token); + } + + if in_refresh_failure_backoff(server_name) { + debug!("Skipping token refresh for MCP server '{server_name}': recent attempt failed"); + return None; + } + + let lock = refresh_lock(server_name); + let _guard = lock.lock().await; + + // A concurrent caller may have refreshed while we waited for the lock. + let tokens = load_oauth_tokens(&key)?; + if Utc::now().timestamp() < tokens.expires_at { + return Some(tokens.access_token); + } + + if in_refresh_failure_backoff(server_name) { + debug!("Skipping token refresh for MCP server '{server_name}': recent attempt failed"); + return None; + } + + match refresh_mcp_token(server_name, &key, &tokens).await { + Ok(access_token) => Some(access_token), + Err(e) => { + note_refresh_failure(server_name); + warn!( + "Failed to refresh OAuth token for MCP server '{server_name}'. \ + Run `.mcp auth {server_name}` to re-authenticate." + ); + debug!( + "Token refresh error for MCP server '{server_name}': {}", + redact_refresh_error(&e) + ); + None + } + } +} + +async fn refresh_mcp_token(server_name: &str, key: &str, tokens: &OAuthTokens) -> Result { + if tokens.refresh_token.is_none() { + return Err(anyhow!("no refresh token stored")); + } + + let reg = + load_registration(server_name).ok_or_else(|| anyhow!("no cached client registration"))?; + let token_url = reg.token_url.ok_or_else(|| { + anyhow!("cached registration has no token URL (saved by an older version)") + })?; + let resource = reg.resource.ok_or_else(|| { + anyhow!("cached registration has no resource (saved by an older version)") + })?; + + let provider = McpOAuthProvider { + client_id: reg.client_id, + authorize_url: String::new(), + token_url, + scopes: String::new(), + fixed_redirect: String::new(), + resource, + }; + + let client = Client::builder().timeout(REFRESH_HTTP_TIMEOUT).build()?; + let refreshed = refresh_oauth_token(&client, &provider, key, tokens).await?; + Ok(refreshed.access_token) +} + +fn refresh_lock(server_name: &str) -> Arc> { + static LOCKS: OnceLock>>>> = + OnceLock::new(); + LOCKS + .get_or_init(Default::default) + .lock() + .entry(server_name.to_string()) + .or_default() + .clone() +} + +fn refresh_failures() -> &'static parking_lot::Mutex> { + static FAILURES: OnceLock>> = OnceLock::new(); + FAILURES.get_or_init(Default::default) +} + +fn note_refresh_failure(server_name: &str) { + refresh_failures() + .lock() + .insert(server_name.to_string(), Instant::now()); +} + +fn in_refresh_failure_backoff(server_name: &str) -> bool { + refresh_failures() + .lock() + .get(server_name) + .is_some_and(|failed_at| failed_at.elapsed() < REFRESH_FAILURE_BACKOFF) +} + +/// Refresh errors may embed the token endpoint's JSON response, which can +/// contain live tokens; strip everything from the first `{` before logging. +fn redact_refresh_error(e: &anyhow::Error) -> String { + let msg = e.to_string(); + match msg.find('{') { + Some(idx) => format!("{}", &msg[..idx]), + None => msg, } } @@ -187,7 +322,13 @@ fn load_registration(server_name: &str) -> Option { serde_json::from_str(&content).ok() } -fn save_registration(server_name: &str, client_id: &str, redirect_uri: &str) -> Result<()> { +fn save_registration( + server_name: &str, + client_id: &str, + redirect_uri: &str, + token_url: &str, + resource: &str, +) -> Result<()> { let dir = paths::oauth_tokens_dir(); fs::create_dir_all(&dir)?; @@ -195,6 +336,8 @@ fn save_registration(server_name: &str, client_id: &str, redirect_uri: &str) -> let reg = McpRegistration { client_id: client_id.to_string(), redirect_uri: Some(redirect_uri.to_string()), + token_url: Some(token_url.to_string()), + resource: Some(resource.to_string()), }; fs::write(path, serde_json::to_string_pretty(®)?)?; @@ -244,7 +387,12 @@ async fn register_client(endpoint: &str, redirect_uri: &str) -> Result { response["client_id"] .as_str() - .ok_or_else(|| anyhow!("Missing client_id in registration response: {response}")) + .ok_or_else(|| { + anyhow!( + "Missing client_id in registration response (keys: {})", + token_response_keys(&response) + ) + }) .map(|s| s.to_string()) } @@ -426,11 +574,31 @@ mod tests { use crate::utils::get_env_name; use serial_test::serial; use std::{ - env, fs, + env, + ffi::OsString, + fs, + path::PathBuf, time::{self, SystemTime}, }; fn with_temp_cache(f: F) { + struct Restore { + key: String, + prev: Option, + root: PathBuf, + } + impl Drop for Restore { + fn drop(&mut self) { + unsafe { + match self.prev.take() { + Some(v) => env::set_var(&self.key, v), + None => env::remove_var(&self.key), + } + } + let _ = fs::remove_dir_all(&self.root); + } + } + let unique = SystemTime::now() .duration_since(time::UNIX_EPOCH) .unwrap() @@ -442,14 +610,12 @@ mod tests { unsafe { env::set_var(&env_key, &root); } + let _restore = Restore { + key: env_key, + prev, + root, + }; f(); - unsafe { - match prev { - Some(v) => env::set_var(&env_key, v), - None => env::remove_var(&env_key), - } - } - let _ = fs::remove_dir_all(&root); } #[test] @@ -685,12 +851,19 @@ mod tests { "notion", "client-xyz-123", "http://127.0.0.1:49152/callback", + "https://as.example/token", + "https://mcp.example/mcp", ) .unwrap(); - let loaded = load_registration("notion"); + let loaded = load_registration("notion").unwrap(); - assert_eq!(loaded.unwrap().client_id, "client-xyz-123"); + assert_eq!(loaded.client_id, "client-xyz-123"); + assert_eq!( + loaded.token_url.as_deref(), + Some("https://as.example/token") + ); + assert_eq!(loaded.resource.as_deref(), Some("https://mcp.example/mcp")); }); } @@ -708,8 +881,22 @@ mod tests { #[serial] fn registration_second_save_overwrites_first() { with_temp_cache(|| { - save_registration("github", "first-id", "http://127.0.0.1:49152/callback").unwrap(); - save_registration("github", "second-id", "http://127.0.0.1:49153/callback").unwrap(); + save_registration( + "github", + "first-id", + "http://127.0.0.1:49152/callback", + "https://as.example/token", + "https://mcp.example/mcp", + ) + .unwrap(); + save_registration( + "github", + "second-id", + "http://127.0.0.1:49153/callback", + "https://as.example/token", + "https://mcp.example/mcp", + ) + .unwrap(); let loaded = load_registration("github").unwrap(); @@ -737,6 +924,8 @@ mod tests { assert_eq!(loaded.client_id, "legacy-id"); assert_eq!(loaded.redirect_uri, None); + assert_eq!(loaded.token_url, None); + assert_eq!(loaded.resource, None); }); } @@ -744,7 +933,14 @@ mod tests { #[serial] fn save_registration_persists_redirect_uri() { with_temp_cache(|| { - save_registration("aws", "client-abc", "http://127.0.0.1:49152/callback").unwrap(); + save_registration( + "aws", + "client-abc", + "http://127.0.0.1:49152/callback", + "https://as.example/token", + "https://mcp.example/mcp", + ) + .unwrap(); let loaded = load_registration("aws").unwrap(); @@ -756,6 +952,82 @@ mod tests { }); } + #[test] + fn mcp_registration_deserializes_without_new_fields_and_roundtrips() { + let old: McpRegistration = serde_json::from_str(r#"{"client_id":"legacy-id"}"#).unwrap(); + + assert_eq!(old.client_id, "legacy-id"); + assert_eq!(old.token_url, None); + assert_eq!(old.resource, None); + + let full = McpRegistration { + client_id: "client-abc".into(), + redirect_uri: Some("http://127.0.0.1:49152/callback".into()), + token_url: Some("https://as.example/token".into()), + resource: Some("https://mcp.example/mcp".into()), + }; + let json = serde_json::to_string(&full).unwrap(); + let back: McpRegistration = serde_json::from_str(&json).unwrap(); + + assert_eq!(back.token_url.as_deref(), Some("https://as.example/token")); + assert_eq!(back.resource.as_deref(), Some("https://mcp.example/mcp")); + } + + #[test] + #[serial] + fn expired_token_with_old_format_registration_returns_none() { + with_temp_cache(|| { + let dir = paths::oauth_tokens_dir(); + fs::create_dir_all(&dir).unwrap(); + fs::write( + paths::token_file("mcp_legacyref"), + r#"{"access_token":"stale","refresh_token":"refresh-abc","expires_at":0}"#, + ) + .unwrap(); + fs::write( + dir.join("mcp_legacyref_registration.json"), + r#"{"client_id":"legacy-id"}"#, + ) + .unwrap(); + + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let token = rt.block_on(load_or_refresh_mcp_token("legacyref")); + + assert_eq!(token, None); + }); + } + + #[test] + fn refresh_failure_backoff_memoizes_per_server() { + assert!(!in_refresh_failure_backoff("backoff-test-server")); + + note_refresh_failure("backoff-test-server"); + + assert!(in_refresh_failure_backoff("backoff-test-server")); + assert!(!in_refresh_failure_backoff("backoff-other-server")); + } + + #[test] + fn redact_refresh_error_strips_response_body() { + let with_body = anyhow!( + "Missing access_token in refresh response: {}", + r#"{"access_token":"live-secret"}"# + ); + let without_body = anyhow!("no refresh token stored"); + + assert_eq!( + redact_refresh_error(&with_body), + "Missing access_token in refresh response: " + ); + assert_eq!( + redact_refresh_error(&without_body), + "no refresh token stored" + ); + } + #[test] fn cached_redirect_port_matches() { let port = cached_redirect_port("http://127.0.0.1:49152/callback", "127.0.0.1", None);