Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ae96a7e031
|
+1
-3
@@ -5,6 +5,4 @@
|
|||||||
.idea/
|
.idea/
|
||||||
/coyote.iml
|
/coyote.iml
|
||||||
/.idea/
|
/.idea/
|
||||||
.coyote/**
|
.coyote
|
||||||
.sisyphus/**
|
|
||||||
.coyote-project.json
|
|
||||||
|
|||||||
@@ -33,11 +33,7 @@ source "$LLM_PROMPT_UTILS_FILE"
|
|||||||
|
|
||||||
# shellcheck disable=SC2154
|
# shellcheck disable=SC2154
|
||||||
main() {
|
main() {
|
||||||
# Command substitution strips *all* trailing newlines and `jq -r` appends one
|
argc_contents="$(jq -r '.content' <<< "$LLM_TOOL_RAW_JSON")"
|
||||||
# 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
|
||||||
@@ -45,11 +41,7 @@ main() {
|
|||||||
exit 1
|
exit 1
|
||||||
fi
|
fi
|
||||||
|
|
||||||
# Same sentinel guard on the patched result, otherwise the trailing newline
|
new_contents="$(patch_file "$argc_path" <(printf "%s" "$argc_contents"))"
|
||||||
# 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?"
|
||||||
|
|||||||
@@ -15,12 +15,7 @@ source "$LLM_PROMPT_UTILS_FILE"
|
|||||||
|
|
||||||
# shellcheck disable=SC2154
|
# shellcheck disable=SC2154
|
||||||
main() {
|
main() {
|
||||||
# Command substitution strips *all* trailing newlines and `jq -r` appends one
|
argc_contents="$(jq -r '.content' <<< "$LLM_TOOL_RAW_JSON")"
|
||||||
# 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
|
||||||
|
|||||||
@@ -841,11 +841,6 @@
|
|||||||
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
|
||||||
|
|||||||
+35
-318
@@ -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, Error, Result, anyhow, bail};
|
use anyhow::{Context, 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, StatusCode};
|
use reqwest::{Client as ReqwestClient, RequestBuilder};
|
||||||
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,10 +16,6 @@ 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;
|
||||||
|
|
||||||
@@ -201,20 +197,10 @@ 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 => {
|
OAuthFlow::ClientCredentials => run_client_credentials_flow(provider, client_name).await,
|
||||||
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,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -315,20 +301,12 @@ 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(|| {
|
.ok_or_else(|| anyhow!("Missing access_token in response: {response}"))?
|
||||||
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"].as_i64().ok_or_else(|| {
|
let expires_in = response["expires_in"]
|
||||||
anyhow!(
|
.as_i64()
|
||||||
"Missing expires_in in response (keys: {})",
|
.ok_or_else(|| anyhow!("Missing expires_in in response: {response}"))?;
|
||||||
token_response_keys(&response)
|
|
||||||
)
|
|
||||||
})?;
|
|
||||||
|
|
||||||
let expires_at = Utc::now().timestamp() + expires_in;
|
let expires_at = Utc::now().timestamp() + expires_in;
|
||||||
|
|
||||||
@@ -356,9 +334,7 @@ async fn run_client_credentials_flow(
|
|||||||
provider: &dyn OAuthProvider,
|
provider: &dyn OAuthProvider,
|
||||||
client_name: &str,
|
client_name: &str,
|
||||||
) -> Result<()> {
|
) -> Result<()> {
|
||||||
let client = ReqwestClient::builder()
|
let client = ReqwestClient::new();
|
||||||
.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"),
|
||||||
@@ -373,19 +349,11 @@ 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(|| {
|
.ok_or_else(|| anyhow!("Missing access_token in client_credentials response: {response}"))?
|
||||||
anyhow!(
|
|
||||||
"Missing access_token in client_credentials response (keys: {})",
|
|
||||||
token_response_keys(&response)
|
|
||||||
)
|
|
||||||
})?
|
|
||||||
.to_string();
|
.to_string();
|
||||||
let expires_in = response["expires_in"].as_i64().ok_or_else(|| {
|
let expires_in = response["expires_in"]
|
||||||
anyhow!(
|
.as_i64()
|
||||||
"Missing expires_in in client_credentials response (keys: {})",
|
.ok_or_else(|| anyhow!("Missing expires_in in client_credentials response: {response}"))?;
|
||||||
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 {
|
||||||
@@ -395,6 +363,11 @@ 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(())
|
||||||
}
|
}
|
||||||
@@ -444,28 +417,19 @@ 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!(
|
anyhow!("Missing device_code in device authorization response: {device_response}")
|
||||||
"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!(
|
anyhow!("Missing user_code in device authorization response: {device_response}")
|
||||||
"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!(
|
anyhow!("Missing verification_uri in device authorization response: {device_response}")
|
||||||
"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"]
|
||||||
@@ -587,76 +551,10 @@ 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)?;
|
||||||
// Write-then-rename so a crash mid-write never truncates the live token file.
|
fs::write(path, json)?;
|
||||||
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,
|
||||||
@@ -679,23 +577,19 @@ pub async fn refresh_oauth_token(
|
|||||||
],
|
],
|
||||||
);
|
);
|
||||||
|
|
||||||
let (status, response) = tokio::time::timeout(TOKEN_ENDPOINT_TIMEOUT, async {
|
let response: Value = request.send().await?.json().await?;
|
||||||
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, refresh_token, expires_in) =
|
let access_token = response["access_token"]
|
||||||
parse_refresh_response(status, &response, tokens.refresh_token.as_deref())?;
|
.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 expires_at = Utc::now().timestamp() + expires_in;
|
let expires_at = Utc::now().timestamp() + expires_in;
|
||||||
|
|
||||||
@@ -715,20 +609,6 @@ 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,
|
||||||
@@ -744,38 +624,18 @@ 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).ok_or_else(|| {
|
load_oauth_tokens(client_name)
|
||||||
anyhow!("Token file missing after client_credentials refresh")
|
.ok_or_else(|| anyhow!("Token file missing after client_credentials refresh"))?
|
||||||
})?
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
tokens
|
tokens
|
||||||
}
|
|
||||||
} else {
|
|
||||||
tokens
|
|
||||||
};
|
};
|
||||||
|
|
||||||
set_access_token(
|
set_access_token(
|
||||||
@@ -1026,54 +886,11 @@ 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 {
|
||||||
@@ -1651,104 +1468,4 @@ 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"));
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -103,7 +103,7 @@ impl McpFactory {
|
|||||||
}
|
}
|
||||||
|
|
||||||
let bearer_token = if spec.is_remote() {
|
let bearer_token = if spec.is_remote() {
|
||||||
oauth::load_or_refresh_mcp_token(name).await
|
oauth::load_valid_mcp_token(name)
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
|||||||
+1
-1
@@ -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_or_refresh_mcp_token(&id).await
|
oauth::load_valid_mcp_token(&id)
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
|||||||
+27
-299
@@ -1,25 +1,15 @@
|
|||||||
use crate::client::oauth::{
|
use crate::client::oauth::{OAuthProvider, TokenRequestFormat, load_oauth_tokens, run_oauth_flow};
|
||||||
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::{debug, warn};
|
use log::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)]
|
||||||
@@ -44,10 +34,6 @@ 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 {
|
||||||
@@ -138,19 +124,8 @@ pub async fn run_mcp_oauth_flow(
|
|||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
|
||||||
let (client_id, redirect_uri) = if let Some((client_id, redirect_uri)) = cached_reuse {
|
let (client_id, redirect_uri) = if let Some(reused) = cached_reuse {
|
||||||
// Re-save so registrations cached before token_url/resource were
|
reused
|
||||||
// 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)?;
|
||||||
@@ -162,7 +137,10 @@ 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) => id,
|
Ok(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:")
|
||||||
@@ -175,18 +153,6 @@ 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)
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -202,115 +168,14 @@ 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 async fn load_or_refresh_mcp_token(server_name: &str) -> Option<String> {
|
pub fn load_valid_mcp_token(server_name: &str) -> Option<String> {
|
||||||
let key = mcp_token_key(server_name);
|
let tokens = load_oauth_tokens(&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 {
|
||||||
return Some(tokens.access_token);
|
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 {
|
||||||
format!("mcp_{server_name}")
|
format!("mcp_{server_name}")
|
||||||
@@ -322,13 +187,7 @@ fn load_registration(server_name: &str) -> Option<McpRegistration> {
|
|||||||
serde_json::from_str(&content).ok()
|
serde_json::from_str(&content).ok()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn save_registration(
|
fn save_registration(server_name: &str, client_id: &str, redirect_uri: &str) -> Result<()> {
|
||||||
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)?;
|
||||||
|
|
||||||
@@ -336,8 +195,6 @@ fn save_registration(
|
|||||||
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(®)?)?;
|
fs::write(path, serde_json::to_string_pretty(®)?)?;
|
||||||
@@ -387,12 +244,7 @@ async fn register_client(endpoint: &str, redirect_uri: &str) -> Result<String> {
|
|||||||
|
|
||||||
response["client_id"]
|
response["client_id"]
|
||||||
.as_str()
|
.as_str()
|
||||||
.ok_or_else(|| {
|
.ok_or_else(|| anyhow!("Missing client_id in registration response: {response}"))
|
||||||
anyhow!(
|
|
||||||
"Missing client_id in registration response (keys: {})",
|
|
||||||
token_response_keys(&response)
|
|
||||||
)
|
|
||||||
})
|
|
||||||
.map(|s| s.to_string())
|
.map(|s| s.to_string())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -574,31 +426,11 @@ 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,
|
env, fs,
|
||||||
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()
|
||||||
@@ -610,12 +442,14 @@ 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]
|
||||||
@@ -851,19 +685,12 @@ 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").unwrap();
|
let loaded = load_registration("notion");
|
||||||
|
|
||||||
assert_eq!(loaded.client_id, "client-xyz-123");
|
assert_eq!(loaded.unwrap().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"));
|
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -881,22 +708,8 @@ mod tests {
|
|||||||
#[serial]
|
#[serial]
|
||||||
fn registration_second_save_overwrites_first() {
|
fn registration_second_save_overwrites_first() {
|
||||||
with_temp_cache(|| {
|
with_temp_cache(|| {
|
||||||
save_registration(
|
save_registration("github", "first-id", "http://127.0.0.1:49152/callback").unwrap();
|
||||||
"github",
|
save_registration("github", "second-id", "http://127.0.0.1:49153/callback").unwrap();
|
||||||
"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();
|
||||||
|
|
||||||
@@ -924,8 +737,6 @@ 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);
|
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -933,14 +744,7 @@ mod tests {
|
|||||||
#[serial]
|
#[serial]
|
||||||
fn save_registration_persists_redirect_uri() {
|
fn save_registration_persists_redirect_uri() {
|
||||||
with_temp_cache(|| {
|
with_temp_cache(|| {
|
||||||
save_registration(
|
save_registration("aws", "client-abc", "http://127.0.0.1:49152/callback").unwrap();
|
||||||
"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();
|
||||||
|
|
||||||
@@ -952,82 +756,6 @@ 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
@@ -3044,7 +3044,7 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn reciprocal_rank_fusion_empty_lists() {
|
fn reciprocal_rank_fusion_empty_lists() {
|
||||||
let result = reciprocal_rank_fusion(vec![], vec![], 5);
|
let result = super::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 = reciprocal_rank_fusion(
|
let result = super::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 = reciprocal_rank_fusion(vec![docs], vec![1.0], 3);
|
let result = super::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 = reciprocal_rank_fusion(
|
let result = super::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,
|
||||||
|
|||||||
+2
-34
@@ -39,10 +39,8 @@ 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> =
|
static BOLD_AST_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 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> =
|
||||||
@@ -2421,36 +2419,6 @@ 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();
|
||||||
|
|||||||
Reference in New Issue
Block a user