feat: identity-aware rejected-token marker for LLM OAuth cache

distrust_access_token compare-and-invalidates the in-memory entry only
when the cached token equals the rejected one, so a concurrent refresh
is never clobbered. is_valid_access_token and both expiry checks in
prepare_oauth_access_token treat marked tokens as expired, forcing a
refresh of provider-rejected tokens that are still locally unexpired.
The marker is cleared after every completed refresh, including ones
that return the same token.
This commit is contained in:
2026-08-14 13:07:51 -06:00
parent f3d59ade11
commit 684f19250a
2 changed files with 204 additions and 4 deletions
+117 -1
View File
@@ -2,6 +2,7 @@ use anyhow::{Result, anyhow};
use chrono::Utc;
use indexmap::IndexMap;
use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::LazyLock;
type AccessTokenEntry = (String, i64, Option<String>);
@@ -9,6 +10,12 @@ type AccessTokenEntry = (String, i64, Option<String>);
static ACCESS_TOKENS: LazyLock<RwLock<IndexMap<String, AccessTokenEntry>>> =
LazyLock::new(|| RwLock::new(IndexMap::new()));
/// Tokens a provider rejected (401) despite being locally unexpired.
/// Maps client name → the exact rejected token so a concurrently-refreshed
/// different token is never distrusted by mistake.
static REJECTED_TOKENS: LazyLock<RwLock<HashMap<String, String>>> =
LazyLock::new(|| RwLock::new(HashMap::new()));
pub fn get_access_token(client_name: &str) -> Result<String> {
ACCESS_TOKENS
.read()
@@ -30,7 +37,7 @@ pub fn is_valid_access_token(client_name: &str) -> bool {
Some(v) => v,
None => return false,
};
!token.is_empty() && Utc::now().timestamp() < *expires_at
!token.is_empty() && Utc::now().timestamp() < *expires_at && !is_rejected(client_name, token)
}
pub fn set_access_token(
@@ -45,3 +52,112 @@ pub fn set_access_token(
entry.1 = expires_at;
entry.2 = account_id;
}
/// Compare-and-invalidate a provider-rejected token.
///
/// Only if the currently-cached token EQUALS `rejected` is the cache entry
/// removed and the rejection marker recorded; a concurrently-refreshed
/// different token is left untouched and no marker is set.
///
/// Returns true if a cache entry existed for this client at all (whether or
/// not it matched `rejected`) — i.e. the client is token-authed and a retry
/// after refresh is worthwhile. Returns false when there is no entry
/// (API-key clients).
#[allow(dead_code)] // Called by the 401-retry path once it lands.
pub fn distrust_access_token(client_name: &str, rejected: &str) -> bool {
let mut access_tokens = ACCESS_TOKENS.write();
let (token, _, _) = match access_tokens.get(client_name) {
Some(v) => v,
None => return false,
};
if token == rejected {
access_tokens.shift_remove(client_name);
REJECTED_TOKENS
.write()
.insert(client_name.to_string(), rejected.to_string());
}
true
}
pub fn is_rejected(client_name: &str, token: &str) -> bool {
REJECTED_TOKENS
.read()
.get(client_name)
.is_some_and(|rejected| rejected == token)
}
pub fn clear_rejected(client_name: &str) {
REJECTED_TOKENS.write().remove(client_name);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn distrust_removes_matching_token_and_sets_marker() {
let client = "distrust-match-test";
set_access_token(client, "at-1".into(), Utc::now().timestamp() + 3600, None);
assert!(distrust_access_token(client, "at-1"));
assert!(get_access_token(client).is_err(), "cache entry not removed");
assert!(is_rejected(client, "at-1"), "marker not set");
}
#[test]
fn distrust_keeps_differing_token_and_skips_marker() {
let client = "distrust-differ-test";
set_access_token(client, "at-new".into(), Utc::now().timestamp() + 3600, None);
assert!(distrust_access_token(client, "at-old"));
assert_eq!(get_access_token(client).unwrap(), "at-new");
assert!(!is_rejected(client, "at-old"), "marker set for stale token");
}
#[test]
fn distrust_returns_false_without_cache_entry() {
let client = "distrust-missing-test";
assert!(!distrust_access_token(client, "at-1"));
assert!(!is_rejected(client, "at-1"));
}
#[test]
fn is_valid_access_token_false_for_rejected_token() {
let client = "rejected-valid-test";
set_access_token(client, "at-1".into(), Utc::now().timestamp() + 3600, None);
assert!(is_valid_access_token(client));
distrust_access_token(client, "at-1");
// A concurrent in-flight prepare re-caches the rejected file token
// between mark and refresh; it must still be treated as invalid.
set_access_token(client, "at-1".into(), Utc::now().timestamp() + 3600, None);
assert!(!is_valid_access_token(client));
}
#[test]
fn clear_rejected_clears_marker_and_clients_are_isolated() {
let client_a = "rejected-isolation-a";
let client_b = "rejected-isolation-b";
set_access_token(client_a, "at-1".into(), Utc::now().timestamp() + 3600, None);
distrust_access_token(client_a, "at-1");
assert!(is_rejected(client_a, "at-1"));
assert!(
!is_rejected(client_b, "at-1"),
"marker leaked across clients"
);
clear_rejected(client_b);
assert!(
is_rejected(client_a, "at-1"),
"wrong client's marker cleared"
);
clear_rejected(client_a);
assert!(!is_rejected(client_a, "at-1"));
}
}