Compare commits

...
Author SHA1 Message Date
Dark-Alex-17 d31110cd67 fix: allow nested italics inside bold spans in markdown renderer
CI / All (ubuntu-latest) (push) Failing after 29s
CI / All (macos-latest) (push) Has been cancelled
CI / All (windows-latest) (push) Has been cancelled
2026-08-13 16:18:50 -06:00
Dark-Alex-17 0f35e03a85 fix: properly handle OAuth refreshes
CI / All (ubuntu-latest) (push) Failing after 29s
CI / All (macos-latest) (push) Has been cancelled
CI / All (windows-latest) (push) Has been cancelled
2026-08-13 16:00:42 -06:00
Dark-Alex-17 c2b0c120d7 chore: Added grok4.6 to models.yaml
CI / All (ubuntu-latest) (push) Failing after 29s
CI / All (macos-latest) (push) Has been cancelled
CI / All (windows-latest) (push) Has been cancelled
2026-08-13 14:20:35 -06:00
Dark-Alex-17 65c9be36b2 fix: correct newline removal from fs_write and fs_patch 2026-08-13 14:18:45 -06:00
11 changed files with 689 additions and 85 deletions
+3 -1
View File
@@ -5,4 +5,6 @@
.idea/ .idea/
/coyote.iml /coyote.iml
/.idea/ /.idea/
.coyote .coyote/**
.sisyphus/**
.coyote-project.json
+10 -2
View File
@@ -33,7 +33,11 @@ source "$LLM_PROMPT_UTILS_FILE"
# shellcheck disable=SC2154 # shellcheck disable=SC2154
main() { main() {
argc_contents="$(jq -r '.content' <<< "$LLM_TOOL_RAW_JSON")" # Command substitution strips *all* trailing newlines and `jq -r` appends one
# of its own, so read with `-j` and pin the real end of the content with a
# sentinel that is removed afterwards.
argc_contents="$(jq -j '.content' <<< "$LLM_TOOL_RAW_JSON"; printf x)"
argc_contents="${argc_contents%x}"
argc_path="$(jq -r '.path' <<< "$LLM_TOOL_RAW_JSON")" argc_path="$(jq -r '.path' <<< "$LLM_TOOL_RAW_JSON")"
if [[ ! -f "$argc_path" ]]; then if [[ ! -f "$argc_path" ]]; then
@@ -41,7 +45,11 @@ main() {
exit 1 exit 1
fi fi
new_contents="$(patch_file "$argc_path" <(printf "%s" "$argc_contents"))" # Same sentinel guard on the patched result, otherwise the trailing newline
# is stripped again on the way back out. `rc` preserves patch_file's exit
# status so a failure still aborts under `set -e`.
new_contents="$(patch_file "$argc_path" <(printf "%s" "$argc_contents"); rc=$?; printf x; exit "$rc")"
new_contents="${new_contents%x}"
printf "%s" "$new_contents" | git diff --no-index "$argc_path" - || true printf "%s" "$new_contents" | git diff --no-index "$argc_path" - || true
guard_operation "Apply changes?" guard_operation "Apply changes?"
+6 -1
View File
@@ -15,7 +15,12 @@ source "$LLM_PROMPT_UTILS_FILE"
# shellcheck disable=SC2154 # shellcheck disable=SC2154
main() { main() {
argc_contents="$(jq -r '.content' <<< "$LLM_TOOL_RAW_JSON")" # Command substitution strips *all* trailing newlines and `jq -r` appends one
# of its own, so read with `-j` and pin the real end of the content with a
# sentinel that is removed afterwards. Without this every written file loses
# its final newline, which breaks formatters such as `cargo fmt --check`.
argc_contents="$(jq -j '.content' <<< "$LLM_TOOL_RAW_JSON"; printf x)"
argc_contents="${argc_contents%x}"
argc_path="$(jq -r '.path' <<< "$LLM_TOOL_RAW_JSON")" argc_path="$(jq -r '.path' <<< "$LLM_TOOL_RAW_JSON")"
if [[ -f "$argc_path" ]]; then if [[ -f "$argc_path" ]]; then
+5
View File
@@ -841,6 +841,11 @@
referrer: coyote referrer: coyote
echo_pkce_in_token_exchange: true echo_pkce_in_token_exchange: true
models: models:
- name: grok-4.6
input_price: 2
output_price: 6
max_input_tokens: 500000
supports_function_calling: true
- name: grok-4.5 - name: grok-4.5
input_price: 2 input_price: 2
output_price: 6 output_price: 6
+1 -4
View File
@@ -765,10 +765,7 @@ mod tests {
assert_eq!(cli.mcp_add, Some("notion".to_string())); assert_eq!(cli.mcp_add, Some("notion".to_string()));
assert!(matches!(cli.transport, Some(McpTransportArg::Http))); assert!(matches!(cli.transport, Some(McpTransportArg::Http)));
assert_eq!(cli.url, Some("https://mcp.notion.com/mcp".to_string())); assert_eq!(cli.url, Some("https://mcp.notion.com/mcp".to_string()));
assert_eq!( assert_eq!(cli.header, vec!["Authorization: Bearer {{NOTION_TOKEN}}"]);
cli.header,
vec!["Authorization: Bearer {{NOTION_TOKEN}}"]
);
assert!(cli.mcp_command.is_empty()); assert!(cli.mcp_command.is_empty());
} }
+318 -35
View File
@@ -2,13 +2,13 @@ use super::access_token::{is_valid_access_token, set_access_token};
use super::openai_compatible_oauth::OpenAICompatibleOAuthProvider; use super::openai_compatible_oauth::OpenAICompatibleOAuthProvider;
use super::{ClientConfig, ProviderModels}; use super::{ClientConfig, ProviderModels};
use crate::config::paths; use crate::config::paths;
use anyhow::{Context, Result, anyhow, bail}; use anyhow::{Context, Error, Result, anyhow, bail};
use base64::Engine; use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD; use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use chrono::Utc; use chrono::Utc;
use indexmap::IndexMap; use indexmap::IndexMap;
use inquire::Text; use inquire::Text;
use reqwest::{Client as ReqwestClient, RequestBuilder}; use reqwest::{Client as ReqwestClient, RequestBuilder, StatusCode};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::Value; use serde_json::Value;
use sha2::{Digest, Sha256}; use sha2::{Digest, Sha256};
@@ -16,6 +16,10 @@ use std::collections::HashMap;
use std::fs; use std::fs;
use std::io::{BufRead, BufReader, Write}; use std::io::{BufRead, BufReader, Write};
use std::net::TcpListener; use std::net::TcpListener;
use std::path::PathBuf;
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use tokio::sync;
use url::Url; use url::Url;
use uuid::Uuid; use uuid::Uuid;
@@ -197,10 +201,20 @@ pub struct OAuthTokens {
pub account_id: Option<String>, pub account_id: Option<String>,
} }
const TOKEN_ENDPOINT_TIMEOUT: Duration = Duration::from_secs(30);
pub async fn run_oauth_flow(provider: &dyn OAuthProvider, client_name: &str) -> Result<()> { pub async fn run_oauth_flow(provider: &dyn OAuthProvider, client_name: &str) -> Result<()> {
match provider.flow() { match provider.flow() {
OAuthFlow::Pkce => run_pkce_flow(provider, client_name).await, 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, 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"] let access_token = response["access_token"]
.as_str() .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(); .to_string();
let refresh_token = response["refresh_token"].as_str().map(|s| s.to_string()); let refresh_token = response["refresh_token"].as_str().map(|s| s.to_string());
let expires_in = response["expires_in"] let expires_in = response["expires_in"].as_i64().ok_or_else(|| {
.as_i64() anyhow!(
.ok_or_else(|| anyhow!("Missing expires_in in response: {response}"))?; "Missing expires_in in response (keys: {})",
token_response_keys(&response)
)
})?;
let expires_at = Utc::now().timestamp() + expires_in; let expires_at = Utc::now().timestamp() + expires_in;
@@ -334,7 +356,9 @@ async fn run_client_credentials_flow(
provider: &dyn OAuthProvider, provider: &dyn OAuthProvider,
client_name: &str, client_name: &str,
) -> Result<()> { ) -> Result<()> {
let client = ReqwestClient::new(); let client = ReqwestClient::builder()
.timeout(TOKEN_ENDPOINT_TIMEOUT)
.build()?;
let scopes = provider.scopes(); let scopes = provider.scopes();
let mut params: Vec<(&str, &str)> = vec![ let mut params: Vec<(&str, &str)> = vec![
("grant_type", "client_credentials"), ("grant_type", "client_credentials"),
@@ -349,11 +373,19 @@ async fn run_client_credentials_flow(
let access_token = response["access_token"] let access_token = response["access_token"]
.as_str() .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(); .to_string();
let expires_in = response["expires_in"] let expires_in = response["expires_in"].as_i64().ok_or_else(|| {
.as_i64() anyhow!(
.ok_or_else(|| anyhow!("Missing expires_in in client_credentials response: {response}"))?; "Missing expires_in in client_credentials response (keys: {})",
token_response_keys(&response)
)
})?;
let expires_at = Utc::now().timestamp() + expires_in; let expires_at = Utc::now().timestamp() + expires_in;
let tokens = OAuthTokens { let tokens = OAuthTokens {
@@ -363,11 +395,6 @@ async fn run_client_credentials_flow(
account_id: provider.extract_account_id(&response), account_id: provider.extract_account_id(&response),
}; };
save_oauth_tokens(client_name, &tokens)?; save_oauth_tokens(client_name, &tokens)?;
println!(
"Successfully authenticated client '{}' with {} via OAuth (client_credentials). Tokens saved.",
client_name,
provider.provider_name()
);
Ok(()) Ok(())
} }
@@ -417,19 +444,28 @@ async fn run_device_code_flow(provider: &dyn OAuthProvider, client_name: &str) -
let device_code = device_response["device_code"] let device_code = device_response["device_code"]
.as_str() .as_str()
.ok_or_else(|| { .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(); .to_string();
let user_code = device_response["user_code"] let user_code = device_response["user_code"]
.as_str() .as_str()
.ok_or_else(|| { .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(); .to_string();
let verification_uri = device_response["verification_uri"] let verification_uri = device_response["verification_uri"]
.as_str() .as_str()
.ok_or_else(|| { .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(); .to_string();
let verification_uri_complete = device_response["verification_uri_complete"] 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)?; fs::create_dir_all(parent)?;
} }
let json = serde_json::to_string_pretty(tokens)?; 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(()) 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 => "<non-object response>".to_string(),
}
}
fn parse_refresh_response(
status: StatusCode,
response: &Value,
previous_refresh_token: Option<&str>,
) -> Result<(String, Option<String>, 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( pub async fn refresh_oauth_token(
client: &ReqwestClient, client: &ReqwestClient,
provider: &dyn OAuthProvider, 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"] let (access_token, refresh_token, expires_in) =
.as_str() parse_refresh_response(status, &response, tokens.refresh_token.as_deref())?;
.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 expires_at = Utc::now().timestamp() + expires_in; let expires_at = Utc::now().timestamp() + expires_in;
@@ -609,6 +715,20 @@ pub async fn refresh_oauth_token(
Ok(new_tokens) 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<sync::Mutex<()>> {
static GUARDS: OnceLock<parking_lot::Mutex<HashMap<String, Arc<sync::Mutex<()>>>>> =
OnceLock::new();
GUARDS
.get_or_init(Default::default)
.lock()
.entry(client_name.to_string())
.or_default()
.clone()
}
pub async fn prepare_oauth_access_token( pub async fn prepare_oauth_access_token(
client: &ReqwestClient, client: &ReqwestClient,
provider: &dyn OAuthProvider, provider: &dyn OAuthProvider,
@@ -624,18 +744,38 @@ pub async fn prepare_oauth_access_token(
}; };
let tokens = if Utc::now().timestamp() >= tokens.expires_at { let tokens = if Utc::now().timestamp() >= tokens.expires_at {
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() { match provider.flow() {
OAuthFlow::Pkce | OAuthFlow::DeviceCode => { OAuthFlow::Pkce | OAuthFlow::DeviceCode => {
refresh_oauth_token(client, provider, client_name, &tokens).await? refresh_oauth_token(client, provider, client_name, &tokens).await?
} }
OAuthFlow::ClientCredentials => { OAuthFlow::ClientCredentials => {
run_client_credentials_flow(provider, client_name).await?; run_client_credentials_flow(provider, client_name).await?;
load_oauth_tokens(client_name) load_oauth_tokens(client_name).ok_or_else(|| {
.ok_or_else(|| anyhow!("Token file missing after client_credentials refresh"))? anyhow!("Token file missing after client_credentials refresh")
})?
} }
} }
} else { } else {
tokens tokens
}
} else {
tokens
}; };
set_access_token( set_access_token(
@@ -886,11 +1026,54 @@ pub(crate) fn client_config_info(
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use std::ffi::OsString;
use std::path::PathBuf;
use std::str; use std::str;
use std::time::UNIX_EPOCH;
use super::*; use super::*;
use crate::client::openai_compatible::OpenAICompatibleConfig; use crate::client::openai_compatible::OpenAICompatibleConfig;
use crate::client::{ModelData, ProviderModels}; 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: FnOnce()>(f: F) {
struct Restore {
key: String,
prev: Option<OsString>,
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 { fn base_config() -> OAuthConfig {
OAuthConfig { OAuthConfig {
@@ -1468,4 +1651,104 @@ scopes:
"body missing grant_type param: {body}" "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"));
}
} }
+1 -1
View File
@@ -103,7 +103,7 @@ impl McpFactory {
} }
let bearer_token = if spec.is_remote() { let bearer_token = if spec.is_remote() {
oauth::load_valid_mcp_token(name) oauth::load_or_refresh_mcp_token(name).await
} else { } else {
None None
}; };
+1 -1
View File
@@ -327,7 +327,7 @@ impl McpRegistry {
.with_context(|| format!("MCP server not found in config: {id}"))?; .with_context(|| format!("MCP server not found in config: {id}"))?;
let bearer_token = if spec.is_remote() { let bearer_token = if spec.is_remote() {
oauth::load_valid_mcp_token(&id) oauth::load_or_refresh_mcp_token(&id).await
} else { } else {
None None
}; };
+299 -27
View File
@@ -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 crate::config::paths;
use anyhow::{Context, Result, anyhow}; use anyhow::{Context, Result, anyhow};
use chrono::Utc; use chrono::Utc;
use inquire::Text; use inquire::Text;
use log::warn; use log::{debug, warn};
use reqwest::Client; use reqwest::Client;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs; use std::fs;
use std::net::TcpListener; use std::net::TcpListener;
use std::sync::{Arc, OnceLock};
use std::time::{Duration, Instant};
use tokio::sync;
use url::Url; use url::Url;
const REFRESH_HTTP_TIMEOUT: Duration = Duration::from_secs(10);
const REFRESH_FAILURE_BACKOFF: Duration = Duration::from_secs(60);
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
struct ProtectedResourceMetadata { struct ProtectedResourceMetadata {
#[serde(default)] #[serde(default)]
@@ -34,6 +44,10 @@ struct McpRegistration {
client_id: String, client_id: String,
#[serde(default)] #[serde(default)]
redirect_uri: Option<String>, redirect_uri: Option<String>,
#[serde(default)]
token_url: Option<String>,
#[serde(default)]
resource: Option<String>,
} }
struct DiscoveredOAuth { struct DiscoveredOAuth {
@@ -124,8 +138,19 @@ pub async fn run_mcp_oauth_flow(
None None
}; };
let (client_id, redirect_uri) = if let Some(reused) = cached_reuse { let (client_id, redirect_uri) = if let Some((client_id, redirect_uri)) = cached_reuse {
reused // 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 { } else {
let bind_addr = format!("127.0.0.1:{}", callback_port.unwrap_or(0)); let bind_addr = format!("127.0.0.1:{}", callback_port.unwrap_or(0));
let listener = TcpListener::bind(&bind_addr)?; let listener = TcpListener::bind(&bind_addr)?;
@@ -137,10 +162,7 @@ pub async fn run_mcp_oauth_flow(
id.to_string() id.to_string()
} else if let Some(reg_endpoint) = &metadata.registration_endpoint { } else if let Some(reg_endpoint) = &metadata.registration_endpoint {
match register_client(reg_endpoint, &redirect_uri).await { match register_client(reg_endpoint, &redirect_uri).await {
Ok(id) => { Ok(id) => id,
let _ = save_registration(server_name, &id, &redirect_uri);
id
}
Err(e) => { Err(e) => {
warn!("Dynamic client registration failed: {e}. Falling back to manual entry."); warn!("Dynamic client registration failed: {e}. Falling back to manual entry.");
Text::new("Enter the OAuth client ID for this MCP server:") Text::new("Enter the OAuth client ID for this MCP server:")
@@ -153,6 +175,18 @@ pub async fn run_mcp_oauth_flow(
.prompt() .prompt()
.context("Failed to read client ID")? .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) (client_id, redirect_uri)
}; };
@@ -168,13 +202,114 @@ pub async fn run_mcp_oauth_flow(
run_oauth_flow(&provider, &mcp_token_key(server_name)).await run_oauth_flow(&provider, &mcp_token_key(server_name)).await
} }
pub fn load_valid_mcp_token(server_name: &str) -> Option<String> { pub async fn load_or_refresh_mcp_token(server_name: &str) -> Option<String> {
let tokens = load_oauth_tokens(&mcp_token_key(server_name))?; let key = mcp_token_key(server_name);
let tokens = load_oauth_tokens(&key)?;
if Utc::now().timestamp() < tokens.expires_at { if Utc::now().timestamp() < tokens.expires_at {
Some(tokens.access_token) return Some(tokens.access_token);
} else { }
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 None
} }
}
}
async fn refresh_mcp_token(server_name: &str, key: &str, tokens: &OAuthTokens) -> Result<String> {
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<sync::Mutex<()>> {
static LOCKS: OnceLock<parking_lot::Mutex<HashMap<String, Arc<sync::Mutex<()>>>>> =
OnceLock::new();
LOCKS
.get_or_init(Default::default)
.lock()
.entry(server_name.to_string())
.or_default()
.clone()
}
fn refresh_failures() -> &'static parking_lot::Mutex<HashMap<String, Instant>> {
static FAILURES: OnceLock<parking_lot::Mutex<HashMap<String, Instant>>> = 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!("{}<response body redacted>", &msg[..idx]),
None => msg,
}
} }
fn mcp_token_key(server_name: &str) -> String { fn mcp_token_key(server_name: &str) -> String {
@@ -187,7 +322,13 @@ fn load_registration(server_name: &str) -> Option<McpRegistration> {
serde_json::from_str(&content).ok() 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(); let dir = paths::oauth_tokens_dir();
fs::create_dir_all(&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 { let reg = McpRegistration {
client_id: client_id.to_string(), client_id: client_id.to_string(),
redirect_uri: Some(redirect_uri.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(&reg)?)?; fs::write(path, serde_json::to_string_pretty(&reg)?)?;
@@ -244,7 +387,12 @@ async fn register_client(endpoint: &str, redirect_uri: &str) -> Result<String> {
response["client_id"] response["client_id"]
.as_str() .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()) .map(|s| s.to_string())
} }
@@ -426,11 +574,31 @@ mod tests {
use crate::utils::get_env_name; use crate::utils::get_env_name;
use serial_test::serial; use serial_test::serial;
use std::{ use std::{
env, fs, env,
ffi::OsString,
fs,
path::PathBuf,
time::{self, SystemTime}, time::{self, SystemTime},
}; };
fn with_temp_cache<F: FnOnce()>(f: F) { fn with_temp_cache<F: FnOnce()>(f: F) {
struct Restore {
key: String,
prev: Option<OsString>,
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() let unique = SystemTime::now()
.duration_since(time::UNIX_EPOCH) .duration_since(time::UNIX_EPOCH)
.unwrap() .unwrap()
@@ -442,14 +610,12 @@ mod tests {
unsafe { unsafe {
env::set_var(&env_key, &root); env::set_var(&env_key, &root);
} }
let _restore = Restore {
key: env_key,
prev,
root,
};
f(); 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] #[test]
@@ -685,12 +851,19 @@ mod tests {
"notion", "notion",
"client-xyz-123", "client-xyz-123",
"http://127.0.0.1:49152/callback", "http://127.0.0.1:49152/callback",
"https://as.example/token",
"https://mcp.example/mcp",
) )
.unwrap(); .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] #[serial]
fn registration_second_save_overwrites_first() { fn registration_second_save_overwrites_first() {
with_temp_cache(|| { with_temp_cache(|| {
save_registration("github", "first-id", "http://127.0.0.1:49152/callback").unwrap(); save_registration(
save_registration("github", "second-id", "http://127.0.0.1:49153/callback").unwrap(); "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(); let loaded = load_registration("github").unwrap();
@@ -737,6 +924,8 @@ mod tests {
assert_eq!(loaded.client_id, "legacy-id"); assert_eq!(loaded.client_id, "legacy-id");
assert_eq!(loaded.redirect_uri, None); assert_eq!(loaded.redirect_uri, None);
assert_eq!(loaded.token_url, None);
assert_eq!(loaded.resource, None);
}); });
} }
@@ -744,7 +933,14 @@ mod tests {
#[serial] #[serial]
fn save_registration_persists_redirect_uri() { fn save_registration_persists_redirect_uri() {
with_temp_cache(|| { 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(); 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: <response body redacted>"
);
assert_eq!(
redact_refresh_error(&without_body),
"no refresh token stored"
);
}
#[test] #[test]
fn cached_redirect_port_matches() { fn cached_redirect_port_matches() {
let port = cached_redirect_port("http://127.0.0.1:49152/callback", "127.0.0.1", None); let port = cached_redirect_port("http://127.0.0.1:49152/callback", "127.0.0.1", None);
+4 -4
View File
@@ -3044,7 +3044,7 @@ mod tests {
#[test] #[test]
fn reciprocal_rank_fusion_empty_lists() { fn reciprocal_rank_fusion_empty_lists() {
let result = super::reciprocal_rank_fusion(vec![], vec![], 5); let result = reciprocal_rank_fusion(vec![], vec![], 5);
assert!(result.is_empty(), "empty input should produce empty output"); assert!(result.is_empty(), "empty input should produce empty output");
} }
@@ -3052,7 +3052,7 @@ mod tests {
fn reciprocal_rank_fusion_deduplicates_across_signals() { fn reciprocal_rank_fusion_deduplicates_across_signals() {
let doc_a = DocumentId::new(0, 0); let doc_a = DocumentId::new(0, 0);
let doc_b = DocumentId::new(0, 1); let doc_b = DocumentId::new(0, 1);
let result = super::reciprocal_rank_fusion( let result = reciprocal_rank_fusion(
vec![vec![doc_a, doc_b], vec![doc_a, doc_b]], vec![vec![doc_a, doc_b], vec![doc_a, doc_b]],
vec![1.0, 1.0], vec![1.0, 1.0],
5, 5,
@@ -3069,7 +3069,7 @@ mod tests {
#[test] #[test]
fn reciprocal_rank_fusion_respects_top_k() { fn reciprocal_rank_fusion_respects_top_k() {
let docs: Vec<DocumentId> = (0..10).map(|i| DocumentId::new(0, i)).collect(); let docs: Vec<DocumentId> = (0..10).map(|i| DocumentId::new(0, i)).collect();
let result = super::reciprocal_rank_fusion(vec![docs], vec![1.0], 3); let result = reciprocal_rank_fusion(vec![docs], vec![1.0], 3);
assert_eq!(result.len(), 3, "result should be capped at top_k=3"); assert_eq!(result.len(), 3, "result should be capped at top_k=3");
} }
@@ -3077,7 +3077,7 @@ mod tests {
fn reciprocal_rank_fusion_weights_affect_ranking() { fn reciprocal_rank_fusion_weights_affect_ranking() {
let doc_a = DocumentId::new(0, 0); let doc_a = DocumentId::new(0, 0);
let doc_b = DocumentId::new(0, 1); let doc_b = DocumentId::new(0, 1);
let result = super::reciprocal_rank_fusion( let result = reciprocal_rank_fusion(
vec![vec![doc_a, doc_b], vec![doc_b, doc_a]], vec![vec![doc_a, doc_b], vec![doc_b, doc_a]],
vec![10.0, 1.0], vec![10.0, 1.0],
2, 2,
+34 -2
View File
@@ -39,8 +39,10 @@ static INLINE_CODE_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"`([^`\n]+
static IMAGE_RE: LazyLock<Regex> = static IMAGE_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"!\[([^\]]*)\]\(([^)]+)\)").unwrap()); LazyLock::new(|| Regex::new(r"!\[([^\]]*)\]\(([^)]+)\)").unwrap());
static LINK_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\[([^\]]+)\]\(([^)]+)\)").unwrap()); static LINK_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\[([^\]]+)\]\(([^)]+)\)").unwrap());
static BOLD_AST_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\*\*([^*\n]+)\*\*").unwrap()); static BOLD_AST_RE: LazyLock<Regex> =
static BOLD_US_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"__([^_\n]+)__").unwrap()); LazyLock::new(|| Regex::new(r"\*\*((?:[^*\n]|\*(?!\*))+?)\*\*").unwrap());
static BOLD_US_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"__((?:[^_\n]|_(?!_))+?)__").unwrap());
static ITALIC_AST_RE: LazyLock<Regex> = static ITALIC_AST_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"(?<![*\w])\*(?!\s)([^*\n]+?)(?<!\s)\*(?!\*)").unwrap()); LazyLock::new(|| Regex::new(r"(?<![*\w])\*(?!\s)([^*\n]+?)(?<!\s)\*(?!\*)").unwrap());
static ITALIC_US_RE: LazyLock<Regex> = static ITALIC_US_RE: LazyLock<Regex> =
@@ -2419,6 +2421,36 @@ std::error::Error>> {
assert!(result.contains("\x1b[9m"), "strikethrough SGR: {result:?}"); assert!(result.contains("\x1b[9m"), "strikethrough SGR: {result:?}");
} }
#[test]
fn bold_asterisk_wraps_italic_asterisk() {
let styles = test_styles();
let result = apply_inline("**loud *soft* loud**", &styles);
assert!(
!result.contains("**"),
"outer bold markers stripped: {result:?}"
);
assert!(result.contains("\x1b[1m"), "bold SGR present: {result:?}");
assert!(result.contains("\x1b[3m"), "italic SGR present: {result:?}");
assert!(result.contains("soft"));
}
#[test]
fn bold_underscore_wraps_italic_underscore() {
let styles = test_styles();
let result = apply_inline("__loud _soft_ loud__", &styles);
assert!(
!result.contains("__"),
"outer bold markers stripped: {result:?}"
);
assert!(result.contains("\x1b[1m"), "bold SGR present: {result:?}");
assert!(result.contains("\x1b[3m"), "italic SGR present: {result:?}");
assert!(result.contains("soft"));
}
#[test] #[test]
fn bold_wraps_inline_code() { fn bold_wraps_inline_code() {
let styles = test_styles(); let styles = test_styles();