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/
/coyote.iml
/.idea/
.coyote
.coyote/**
.sisyphus/**
.coyote-project.json
+10 -2
View File
@@ -33,7 +33,11 @@ source "$LLM_PROMPT_UTILS_FILE"
# shellcheck disable=SC2154
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")"
if [[ ! -f "$argc_path" ]]; then
@@ -41,7 +45,11 @@ main() {
exit 1
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
guard_operation "Apply changes?"
+6 -1
View File
@@ -15,7 +15,12 @@ source "$LLM_PROMPT_UTILS_FILE"
# shellcheck disable=SC2154
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")"
if [[ -f "$argc_path" ]]; then
+5
View File
@@ -841,6 +841,11 @@
referrer: coyote
echo_pkce_in_token_exchange: true
models:
- name: grok-4.6
input_price: 2
output_price: 6
max_input_tokens: 500000
supports_function_calling: true
- name: grok-4.5
input_price: 2
output_price: 6
+1 -4
View File
@@ -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());
}
+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::{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<String>,
}
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 => "<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(
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<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(
client: &ReqwestClient,
provider: &dyn OAuthProvider,
@@ -624,18 +744,38 @@ pub async fn prepare_oauth_access_token(
};
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() {
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"))?
load_oauth_tokens(client_name).ok_or_else(|| {
anyhow!("Token file missing after client_credentials refresh")
})?
}
}
} else {
tokens
}
} else {
tokens
};
set_access_token(
@@ -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: 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 {
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"));
}
}
+1 -1
View File
@@ -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
};
+1 -1
View File
@@ -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
};
+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 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<String>,
#[serde(default)]
token_url: Option<String>,
#[serde(default)]
resource: Option<String>,
}
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,14 +202,115 @@ 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<String> {
let tokens = load_oauth_tokens(&mcp_token_key(server_name))?;
pub async fn load_or_refresh_mcp_token(server_name: &str) -> Option<String> {
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 {
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<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 {
format!("mcp_{server_name}")
@@ -187,7 +322,13 @@ fn load_registration(server_name: &str) -> Option<McpRegistration> {
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(&reg)?)?;
@@ -244,7 +387,12 @@ async fn register_client(endpoint: &str, redirect_uri: &str) -> Result<String> {
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: 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(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: <response body redacted>"
);
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);
+4 -4
View File
@@ -3044,7 +3044,7 @@ mod tests {
#[test]
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");
}
@@ -3052,7 +3052,7 @@ mod tests {
fn reciprocal_rank_fusion_deduplicates_across_signals() {
let doc_a = DocumentId::new(0, 0);
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![1.0, 1.0],
5,
@@ -3069,7 +3069,7 @@ mod tests {
#[test]
fn reciprocal_rank_fusion_respects_top_k() {
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");
}
@@ -3077,7 +3077,7 @@ mod tests {
fn reciprocal_rank_fusion_weights_affect_ranking() {
let doc_a = DocumentId::new(0, 0);
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![10.0, 1.0],
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> =
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_US_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"__([^_\n]+)__").unwrap());
static BOLD_AST_RE: LazyLock<Regex> =
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> =
LazyLock::new(|| Regex::new(r"(?<![*\w])\*(?!\s)([^*\n]+?)(?<!\s)\*(?!\*)").unwrap());
static ITALIC_US_RE: LazyLock<Regex> =
@@ -2419,6 +2421,36 @@ std::error::Error>> {
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]
fn bold_wraps_inline_code() {
let styles = test_styles();