grok-build-upstream-mirror/crates/codegen/xai-grok-shell/src/auth/manager_tests.rs

4364 lines
160 KiB
Rust
Raw Normal View History

//! Unit tests for [`super::manager::AuthManager`]. Extracted from
//! `manager.rs` so the implementation reads top-to-bottom; wired in
//! via `#[path = "manager_tests.rs"] mod tests;` in manager.rs.
use super::*;
use crate::auth::error::RefreshTokenError;
use std::sync::atomic::{AtomicU32, Ordering};
use std::time::Instant;
fn make_auth(expires_at: Option<DateTime<Utc>>, create_time: DateTime<Utc>) -> GrokAuth {
GrokAuth {
auth_mode: AuthMode::External,
create_time,
user_id: String::new(),
expires_at,
..GrokAuth::test_default()
}
}
#[test]
fn expired_within_5min_buffer() {
let auth = make_auth(Some(Utc::now() + Duration::minutes(4)), Utc::now());
assert!(is_expired(&auth));
}
#[test]
fn fallback_ttl_when_no_expires_at() {
let old = Utc::now() - Duration::days(30) + Duration::minutes(4);
let auth = make_auth(None, old);
assert!(is_expired(&auth));
let recent = Utc::now() - Duration::days(29);
let auth = make_auth(None, recent);
assert!(!is_expired(&auth));
}
#[test]
fn has_usable_disk_token_reads_disk_independent_of_memory() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
assert!(!mgr.has_usable_disk_token());
let valid = make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now());
mgr.persist_and_swap(valid);
mgr.clear_in_memory();
assert!(mgr.current().is_none(), "in-memory cleared");
assert!(
mgr.has_usable_disk_token(),
"a valid token on disk is usable even when in-memory is empty"
);
let expired = make_auth(Some(Utc::now() - Duration::hours(1)), Utc::now());
mgr.persist_and_swap(expired);
mgr.clear_in_memory();
assert!(
!mgr.has_usable_disk_token(),
"an expired token on disk is not usable"
);
}
#[test]
fn has_usable_token_covers_memory_and_disk() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
assert!(!mgr.has_usable_token(), "nothing in memory or on disk");
mgr.hot_swap(make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()));
assert!(!mgr.has_usable_disk_token(), "disk still empty");
assert!(mgr.has_usable_token(), "valid in-memory token is usable");
mgr.persist_and_swap(make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()));
mgr.hot_swap(make_auth(Some(Utc::now() - Duration::hours(1)), Utc::now()));
assert!(mgr.current().is_none(), "in-memory token is expired");
assert!(mgr.has_usable_token(), "fresh disk token keeps it usable");
mgr.persist_and_swap(make_auth(Some(Utc::now() - Duration::hours(1)), Utc::now()));
assert!(
!mgr.has_usable_token(),
"expired in memory and on disk is not usable"
);
}
#[test]
fn auth_scope_uses_oauth2_when_present() {
let cfg = GrokComConfig::default();
// Default config always has oauth2 set to the xAI defaults.
assert_eq!(
cfg.auth_scope(),
format!(
"{}::{}",
crate::auth::config::XAI_OAUTH2_ISSUER,
obfstr::obfstr!("b1a00492-073a-47ea-816f-4c329264a828"),
)
);
}
#[test]
fn legacy_scope_fallback_reads_old_auth_json() {
let dir = tempfile::tempdir().unwrap();
let auth_path = dir.path().join("auth.json");
// Write auth.json with the legacy scope key (as `x setup` copies from
// a machine that was authenticated with an older grok version).
let legacy_auth = make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now());
let mut store = AuthStore::new();
store.insert(LEGACY_SCOPE.to_string(), legacy_auth);
write_auth_json(&auth_path, &store).unwrap();
// AuthManager uses the new OAuth2 scope, but should still find the
// token under the legacy key.
let cfg = GrokComConfig::default();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
let current = mgr.current();
assert!(current.is_some(), "should fall back to legacy scope key");
assert_eq!(current.unwrap().key, "test-key");
}
#[test]
fn new_scope_takes_precedence_over_legacy() {
let dir = tempfile::tempdir().unwrap();
let auth_path = dir.path().join("auth.json");
let legacy_auth = GrokAuth {
key: "legacy-key".into(),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
let new_auth = GrokAuth {
key: "new-key".into(),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
let cfg = GrokComConfig::default();
let scope = cfg.auth_scope();
let mut store = AuthStore::new();
store.insert(LEGACY_SCOPE.to_string(), legacy_auth);
store.insert(scope, new_auth);
write_auth_json(&auth_path, &store).unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
let current = mgr.current().expect("should find auth");
assert_eq!(current.key, "new-key", "new scope should take precedence");
}
// -- Near-expiry (5-minute buffer) behavior ------------------------
/// Regression test: a token within the 5-minute early-invalidation buffer
/// must be invisible to `current()` (returns None) but visible to
/// `expired_auth()` so that callers can attempt a silent refresh.
#[test]
fn near_expiry_token_invisible_to_current_visible_to_expired_auth() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
// Token expires in 3 minutes -- inside the 5-minute buffer.
let near_expiry = GrokAuth {
key: "near-expiry-key".into(),
user_id: "user-1".into(),
email: Some("user@test.com".into()),
refresh_token: Some("rt-valid".into()),
expires_at: Some(Utc::now() + Duration::minutes(3)),
oidc_issuer: Some("https://idp.example.com".into()),
oidc_client_id: Some("client-1".into()),
..GrokAuth::test_default()
};
mgr.hot_swap(near_expiry);
// current() must return None (token is "expired" per buffer)
assert!(
mgr.current().is_none(),
"current() should return None for token within 5-min buffer"
);
// is_expired() must return true
assert!(
mgr.is_expired(),
"is_expired() should be true for token within 5-min buffer"
);
// expired_auth() must return the token so refresh can use it
let expired = mgr.expired_auth();
assert!(
expired.is_some(),
"expired_auth() should return the near-expiry token"
);
assert_eq!(expired.as_ref().unwrap().key, "near-expiry-key");
assert_eq!(
expired.as_ref().unwrap().refresh_token.as_deref(),
Some("rt-valid"),
"refresh_token must be preserved for silent refresh"
);
}
#[tokio::test]
async fn update_preserves_other_scope_entries() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg.clone()));
// Pre-populate with an external auth entry
let external = GrokAuth {
key: "external-key".into(),
auth_mode: AuthMode::External,
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
{
let mut map = AuthStore::new();
map.insert("other-scope".into(), external);
write_auth_json(&dir.path().join("auth.json"), &map).unwrap();
}
// Now update via auth_manager
let new_auth = GrokAuth {
key: "oidc-token".into(),
auth_mode: AuthMode::Oidc,
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
mgr.update(new_auth).await.unwrap();
// Both entries should exist
let store = read_auth_json(&dir.path().join("auth.json")).unwrap();
assert!(store.contains_key("other-scope"));
assert!(store.contains_key(&cfg.auth_scope()));
}
/// Regression: when auth.json contains corrupt JSON, update() must not
/// clobber the file with a single-entry map. Instead it should update
/// in-memory only and leave the file untouched.
#[tokio::test]
async fn update_recovers_from_corrupt_auth_json_by_backing_up_old_file() {
let dir = tempfile::tempdir().unwrap();
let auth_path = dir.path().join("auth.json");
let cfg = GrokComConfig::default();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg.clone()));
let bad_content = b"NOT VALID JSON {{{";
std::fs::write(&auth_path, bad_content).unwrap();
let new_auth = GrokAuth {
key: "fresh-token".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("fresh-rt".into()),
user_id: "fresh-user".into(),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
let result = mgr.update(new_auth).await;
assert!(
result.is_ok(),
"update must succeed and persist after corrupt recovery: {result:?}"
);
let current = mgr.current();
assert_eq!(
current.as_ref().map(|a| a.key.as_str()),
Some("fresh-token")
);
let on_disk_raw = std::fs::read_to_string(&auth_path).unwrap();
assert!(
on_disk_raw.contains("fresh-token"),
"auth.json must contain the new credential after recovery, got: {on_disk_raw}"
);
let on_disk: AuthStore =
serde_json::from_str(&on_disk_raw).expect("auth.json must be valid JSON after recovery");
assert!(on_disk.contains_key(&cfg.auth_scope()));
let mut backup_found = None;
for entry in std::fs::read_dir(dir.path()).unwrap() {
let entry = entry.unwrap();
let name = entry.file_name().to_string_lossy().into_owned();
if name.starts_with("auth.json.corrupt.") {
backup_found = Some(entry.path());
break;
}
}
let backup_path = backup_found.expect("a .corrupt.* backup file must have been created");
let backup_content = std::fs::read_to_string(&backup_path).unwrap();
assert!(
backup_content.contains("NOT VALID JSON"),
"backup must contain the original corrupt content, got: {backup_content}"
);
}
/// Regression test: update() must preserve team fields from the OIDC flow
/// when the proxy `/user` response does not include them.
#[tokio::test]
async fn update_preserves_team_fields_when_proxy_omits_them() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
// Point proxy_base_url to a non-existent server so the /user call
// fails and falls back to the auth-flow values.
let mgr = Arc::new(AuthManager::new(dir.path(), cfg).with_proxy_base_url("http://127.0.0.1:1"));
let team_auth = GrokAuth {
key: "team-token".into(),
auth_mode: AuthMode::Oidc,
principal_type: Some("Team".into()),
principal_id: Some("team-xyz".into()),
team_id: Some("team-xyz".into()),
team_name: None,
team_role: None,
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
let saved = mgr.update(team_auth).await.unwrap();
assert_eq!(
saved.principal_type.as_deref(),
Some("Team"),
"principal_type must survive proxy fallback"
);
assert_eq!(
saved.principal_id.as_deref(),
Some("team-xyz"),
"principal_id must survive proxy fallback"
);
assert_eq!(
saved.team_id.as_deref(),
Some("team-xyz"),
"team_id must survive proxy fallback"
);
// Verify on-disk too
let store = read_auth_json(&dir.path().join("auth.json")).unwrap();
let on_disk = store.values().next().unwrap();
assert_eq!(on_disk.principal_type.as_deref(), Some("Team"));
assert_eq!(on_disk.team_id.as_deref(), Some("team-xyz"));
}
/// Team tokens are stored under the base scope key (same as personal).
/// There is at most one OAuth entry per issuer/client pair.
#[tokio::test]
async fn update_stores_team_token_under_base_scope() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let base_scope = cfg.auth_scope();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg).with_proxy_base_url("http://127.0.0.1:1"));
let team_auth = GrokAuth {
key: "team-token".into(),
auth_mode: AuthMode::Oidc,
principal_type: Some("Team".into()),
principal_id: Some("team-abc".into()),
team_id: Some("team-abc".into()),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
mgr.update(team_auth).await.unwrap();
let store = read_auth_json(&dir.path().join("auth.json")).unwrap();
assert!(
store.contains_key(&base_scope),
"team token must be stored under base scope '{}', found keys: {:?}",
base_scope,
store.keys().collect::<Vec<_>>()
);
assert_eq!(store.get(&base_scope).unwrap().key, "team-token");
}
/// Logging in as personal must evict any existing team token
/// (at most one OAuth session per issuer/client pair).
#[tokio::test]
async fn team_login_then_personal_evicts_team_token() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let base_scope = cfg.auth_scope();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg).with_proxy_base_url("http://127.0.0.1:1"));
// Step 1: login as team
let team_auth = GrokAuth {
key: "team-token".into(),
principal_type: Some("Team".into()),
principal_id: Some("team-abc".into()),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
mgr.update(team_auth).await.unwrap();
// Step 2: login as personal
let personal_auth = GrokAuth {
key: "personal-token".into(),
principal_type: None,
principal_id: None,
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
mgr.update(personal_auth).await.unwrap();
let store = read_auth_json(&dir.path().join("auth.json")).unwrap();
assert_eq!(
store.len(),
1,
"only one OAuth entry should remain, found: {:?}",
store.keys().collect::<Vec<_>>()
);
assert!(store.contains_key(&base_scope));
assert_eq!(store.get(&base_scope).unwrap().key, "personal-token");
}
/// Regression test: clear() must only remove the current scope, not the
/// legacy scope. Previously, logging in with OAuth would also delete the
/// legacy `https://accounts.x.ai/sign-in` entry from auth.json.
#[test]
fn clear_does_not_remove_legacy_scope() {
let dir = tempfile::tempdir().unwrap();
let auth_path = dir.path().join("auth.json");
let legacy_auth = GrokAuth {
key: "legacy-key".into(),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
let oauth_auth = GrokAuth {
key: "oauth-key".into(),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
let cfg = GrokComConfig::default();
let scope = cfg.auth_scope();
let mut store = AuthStore::new();
store.insert(LEGACY_SCOPE.to_string(), legacy_auth);
store.insert(scope, oauth_auth);
write_auth_json(&auth_path, &store).unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
// clear() should only remove the OAuth scope, not legacy
mgr.clear().unwrap();
let on_disk = read_auth_json(&auth_path).unwrap();
assert!(
on_disk.contains_key(LEGACY_SCOPE),
"legacy scope should be preserved after clear()"
);
assert!(
!on_disk.contains_key(&mgr.scope),
"current scope should be removed after clear()"
);
}
#[test]
fn is_data_collection_disabled_matrix() {
// (team_blocked_reasons, coding_data_retention_opt_out, expected)
let cases: &[(&[&str], bool, bool)] = &[
// ZDR team alone
(&["BLOCKED_REASON_NO_LOGS"], false, true),
(&["BLOCKED_REASON_NO_LOGS_MODERATED"], false, true),
// Opt-out alone
(&[], true, true),
// Both
(&["BLOCKED_REASON_NO_LOGS"], true, true),
// Neither
(&[], false, false),
// Unrelated blocked reasons
(
&["BLOCKED_REASON_BILLING", "BLOCKED_REASON_SUSPENDED"],
false,
false,
),
(&["BLOCKED_REASON_BILLING"], true, true),
// ZDR mixed with other reasons
(
&["BLOCKED_REASON_BILLING", "BLOCKED_REASON_NO_LOGS"],
false,
true,
),
];
for (reasons, opt_out, expected) in cases {
let auth = GrokAuth {
team_blocked_reasons: reasons.iter().map(|s| (*s).into()).collect(),
coding_data_retention_opt_out: *opt_out,
..GrokAuth::test_default()
};
assert_eq!(
auth.is_data_collection_disabled(),
*expected,
"reasons={reasons:?} opt_out={opt_out} expected={expected}",
);
}
}
/// Fail-direction contract of the two `AuthManager` collection predicates:
/// `is_data_collection_disabled` fails open on missing credentials (legacy
/// semantics shared by telemetry/sync gates), `allows_data_collection` fails
/// closed (nothing may leave the machine while privacy state is unknown,
/// e.g. after a mid-session `/logout`).
#[test]
fn manager_collection_predicates_fail_directions() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// No credential: disabled=false (fail-open), allows=false (fail-closed).
assert!(!mgr.is_data_collection_disabled());
assert!(
!mgr.allows_data_collection(),
"missing credential must fail closed for collection"
);
// Normal user: both predicates allow collection.
mgr.hot_swap(GrokAuth::test_default());
assert!(!mgr.is_data_collection_disabled());
assert!(mgr.allows_data_collection());
// Opted-out user: both predicates suppress collection.
mgr.hot_swap(GrokAuth {
coding_data_retention_opt_out: true,
..GrokAuth::test_default()
});
assert!(mgr.is_data_collection_disabled());
assert!(!mgr.allows_data_collection());
// Mid-session `/logout`: the fail-closed predicate flips back to
// "no collection" even after a previously permissive credential.
mgr.hot_swap(GrokAuth::test_default());
assert!(mgr.allows_data_collection(), "precondition");
mgr.clear_in_memory();
assert!(
!mgr.allows_data_collection(),
"cleared credentials must close the collection gate"
);
}
// -- token_suffix ----------------------------------------------------------------
#[test]
fn token_suffix_matrix() {
let cases: &[(&str, &str)] = &[
("abcdefghijklmnop", "efghijklmnop"), // takes last 12
("short", "short"), // short unchanged
("", ""), // empty
("123456789012", "123456789012"), // exact 12
];
for (input, expected) in cases {
assert_eq!(token_suffix(input), *expected, "input={input:?}");
}
}
// -- read_disk_auth ----------------------------------------------------------
// -- hot_swap / try_use_disk_token ---------------------------------------
#[test]
fn hot_swap_updates_in_memory_without_disk() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
assert!(mgr.current().is_none());
let auth = GrokAuth {
key: "swapped".into(),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
mgr.hot_swap(auth);
assert_eq!(mgr.current().unwrap().key, "swapped");
// Disk should NOT have the token
assert!(mgr.read_disk_auth().is_none());
}
#[test]
fn try_use_disk_token_accepts_valid_disk_token() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
let valid_disk = GrokAuth {
key: "valid-disk".into(),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
let result = mgr.try_use_disk_token(Some(&valid_disk), RefreshReason::PreRequest);
assert_eq!(result.unwrap().key, "valid-disk");
// Should also hot-swap into memory
assert_eq!(mgr.current().unwrap().key, "valid-disk");
}
#[test]
fn try_use_disk_token_rejects_expired_disk_token() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
let expired_disk = make_auth(Some(Utc::now() - Duration::hours(1)), Utc::now());
assert!(
mgr.try_use_disk_token(Some(&expired_disk), RefreshReason::PreRequest)
.is_none()
);
}
#[test]
fn try_use_disk_token_rejects_same_key_on_server_rejected() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
let auth = GrokAuth {
key: "same-key".into(),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
mgr.hot_swap(auth.clone());
// ServerRejected should not accept a disk token with the same key
assert!(
mgr.try_use_disk_token(Some(&auth), RefreshReason::ServerRejected)
.is_none()
);
}
#[test]
fn try_use_disk_token_accepts_different_key_on_server_rejected() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
let mem_auth = GrokAuth {
key: "old-key".into(),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
mgr.hot_swap(mem_auth);
let disk_auth = GrokAuth {
key: "new-key".into(),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
let result = mgr.try_use_disk_token(Some(&disk_auth), RefreshReason::ServerRejected);
assert_eq!(result.unwrap().key, "new-key");
}
// -- File locking ----------------------------------------------------------
// -- Disk-refresh race simulation ------------------------------------------
/// Simulates the core scenario this PR fixes: an expired in-memory token
/// where another process has already refreshed on disk. The manager should
/// pick up the valid disk token via try_use_disk_token instead of
/// attempting its own refresh.
#[tokio::test]
async fn disk_refresh_wins_over_expired_in_memory() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let scope = cfg.auth_scope();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
// Simulate: in-memory token is expired
let expired = GrokAuth {
key: "expired-key".into(),
refresh_token: Some("old-rt".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(expired);
assert!(mgr.is_expired());
assert!(mgr.current().is_none());
// Simulate: another process wrote a valid token to disk
let fresh_disk = GrokAuth {
key: "fresh-key-from-sibling".into(),
refresh_token: Some("new-rt".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
};
let mut store = AuthStore::new();
store.insert(scope, fresh_disk);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
// Acquire lock + read disk (mirrors flow.rs logic)
let _lock = mgr
.try_lock_auth_file_async(StdDuration::from_secs(1))
.await;
assert!(_lock.is_some());
let disk_auth = mgr.read_disk_auth();
assert!(disk_auth.is_some());
assert!(!is_expired(disk_auth.as_ref().unwrap()));
// try_use_disk_token should accept it and hot-swap
let result = mgr.try_use_disk_token(disk_auth.as_ref(), RefreshReason::PreRequest);
assert_eq!(result.unwrap().key, "fresh-key-from-sibling");
assert_eq!(mgr.current().unwrap().key, "fresh-key-from-sibling");
}
struct CountingRefresher {
call_count: Arc<AtomicU32>,
delay: StdDuration,
}
#[async_trait::async_trait]
impl TokenRefresher for CountingRefresher {
async fn refresh(&self, _reason: RefreshReason) -> crate::auth::refresh::RefreshOutcome {
self.call_count.fetch_add(1, Ordering::SeqCst);
tokio::time::sleep(self.delay).await;
let fresh = GrokAuth {
key: "fresh-token".into(),
expires_at: Some(Utc::now() + Duration::hours(1)),
refresh_token: Some("rt-new".into()),
..GrokAuth::test_default()
};
crate::auth::refresh::RefreshOutcome::Success(Box::new(fresh))
}
}
struct FailingRefresher {
call_count: Arc<AtomicU32>,
}
#[async_trait::async_trait]
impl TokenRefresher for FailingRefresher {
async fn refresh(&self, _reason: RefreshReason) -> crate::auth::refresh::RefreshOutcome {
self.call_count.fetch_add(1, Ordering::SeqCst);
crate::auth::refresh::RefreshOutcome::permanent(
crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected,
None,
)
}
}
/// Record a permanent failure scoped to the auth manager's current (or expired)
/// credential key, mirroring what `refresh_chain` does in production.
fn record_permanent_failure(
auth_manager: &AuthManager,
reason: crate::auth::error::RefreshTokenFailedReason,
) {
let key = auth_manager
.current()
.or_else(|| auth_manager.expired_auth())
.map(|a| a.key)
.unwrap_or_default();
auth_manager.record_permanent_failure(key, reason.into());
}
/// Permanent-failure refresher that reports a specific `tried_key` (the
/// credential it claims to have sent to the IdP), letting tests assert the
/// verdict is keyed on the actually-tried credential.
struct TriedKeyFailRefresher {
tried_key: String,
call_count: Arc<AtomicU32>,
}
#[async_trait::async_trait]
impl TokenRefresher for TriedKeyFailRefresher {
async fn refresh(&self, _reason: RefreshReason) -> crate::auth::refresh::RefreshOutcome {
self.call_count.fetch_add(1, Ordering::SeqCst);
crate::auth::refresh::RefreshOutcome::permanent(
crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected,
Some(self.tried_key.clone()),
)
}
}
/// With `inner == None` but a dead refresh-token on disk, the refresher still
/// exchanges that disk RT. The verdict must be keyed on the
/// credential actually tried (the disk RT), so repeated reactive refreshes
/// short-circuit on it instead of hammering the IdP.
#[tokio::test]
async fn storm_cap_engages_with_empty_inner_and_dead_disk_refresh_token() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let scope = cfg.auth_scope();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
// Disk: an expired token carrying the (dead) refresh_token the OIDC
// refresher resolves. `inner` stays empty.
let dead = GrokAuth {
key: "disk-dead".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-dead".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
};
let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default();
store.insert(scope, dead);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
assert!(mgr.current_or_expired().is_none(), "inner must be empty");
let calls = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(FailingRefresher {
call_count: calls.clone(),
}));
for _ in 0..5 {
let _ = mgr
.refresh_chain(TokenType::OidcSession, RefreshReason::ServerRejected)
.await;
}
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"storm cap must hold the IdP to one call even with empty inner + dead disk RT",
);
}
/// Record/check consistency: in-mem and disk are DIFFERENT stale credentials.
/// The refresher resolves & sends the DISK refresh token, so the verdict must be
/// keyed on THAT — proven by swapping the in-mem bearer afterward and confirming
/// the verdict still caps the storm (a verdict mis-keyed to the in-mem bearer
/// would read absent after the swap and re-hit the IdP). The `tried_key == None`
/// fallback (external-binary flow → `attempted_verdict_key`) is covered by
/// `storm_cap_engages_with_empty_inner_and_dead_disk_refresh_token`.
#[tokio::test]
async fn verdict_not_keyed_on_in_mem_bearer() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let scope = cfg.auth_scope();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
// in-mem: stale bearer K_mem (expired, with RT).
mgr.hot_swap(GrokAuth {
key: "mem-stale".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-mem".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
});
// disk: a DIFFERENT stale credential K_disk (expired, with RT) — what the
// refresher resolves first.
let disk = GrokAuth {
key: "disk-stale".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-disk".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
};
let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default();
store.insert(scope, disk);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
let calls = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(TriedKeyFailRefresher {
tried_key: "disk-stale".into(),
call_count: calls.clone(),
}));
let _ = mgr
.refresh_chain(TokenType::OidcSession, RefreshReason::ServerRejected)
.await;
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"first call hits the IdP once"
);
// Swap the in-mem bearer to yet another stale key: a verdict mis-keyed to
// the old in-mem bearer would now read absent.
mgr.hot_swap(GrokAuth {
key: "mem-stale-2".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-mem-2".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
});
let _ = mgr
.refresh_chain(TokenType::OidcSession, RefreshReason::ServerRejected)
.await;
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"verdict keyed on the tried disk credential must survive an in-mem swap",
);
}
/// Success → persist-failure → transient: a refresh that obtains a fresh token
/// but cannot write it to disk must surface `Transient` AND still swap the
/// in-memory bearer to the fresh token (the "always update in-memory even if the
/// disk write failed" invariant — without it a disk hiccup strands the session).
/// The write is failed deterministically (root-safe) by planting a *directory*
/// at the atomic-write temp path so `open_secure_file` hits `EISDIR`; the
/// auth.json read (file absent) and the file lock still succeed.
#[tokio::test]
async fn refresh_persist_failure_is_transient_but_swaps_in_memory() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// Expired in-mem bearer so the chain proceeds to the IdP (no early return).
mgr.hot_swap(GrokAuth {
key: "stale".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
});
// `write_auth_json_atomic` writes `auth.json.<pid>.tmp` then renames; a
// directory there makes the temp-file open fail with EISDIR (enforced even
// for root), so the persist fails while the read/lock paths are unaffected.
std::fs::create_dir(
dir.path()
.join(format!("auth.json.{}.tmp", std::process::id())),
)
.unwrap();
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: Arc::new(AtomicU32::new(0)),
delay: StdDuration::ZERO,
}));
let err = mgr
.refresh_chain(TokenType::OidcSession, RefreshReason::ServerRejected)
.await
.expect_err("persist failure must surface an error");
assert!(
matches!(err, AuthError::Refresh(RefreshTokenError::Transient(_))),
"persist failure must be transient (retryable), got {err:?}",
);
assert_eq!(
mgr.current().map(|a| a.key),
Some("fresh-token".to_string()),
"in-memory bearer must hold the fresh token despite the failed disk write",
);
}
#[tokio::test]
async fn auth_concurrent_refresh_deduplicates() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let expired = GrokAuth {
key: "expired-key".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-old".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(expired);
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: call_count.clone(),
delay: StdDuration::from_millis(50),
}));
// Spawn 4 concurrent tasks that all call auth().
let mut handles = Vec::new();
for _ in 0..4 {
let m = mgr.clone();
handles.push(tokio::spawn(async move { m.auth().await }));
}
let mut results = Vec::new();
for h in handles {
results.push(h.await.unwrap());
}
// All 4 should succeed with the same fresh token.
for r in &results {
assert_eq!(
r.as_ref().unwrap().key,
"fresh-token",
"all tasks must get the fresh token"
);
}
// The refresher should have been called exactly once.
assert_eq!(
call_count.load(Ordering::SeqCst),
1,
"refresher must be called exactly once despite 4 concurrent callers"
);
}
#[tokio::test]
async fn auth_permanent_failure_stops_retries() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let expired = GrokAuth {
key: "expired-key".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-old".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(expired);
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(FailingRefresher {
call_count: call_count.clone(),
}));
// First auth(): refresher called, refresh_chain records permanent failure.
let err1 = mgr.auth().await.unwrap_err();
assert!(
matches!(err1, AuthError::Refresh(RefreshTokenError::Permanent(_))),
"first call should return PermanentFailure, got: {err1:?}"
);
// Second auth(): permanent failure cached, refresher NOT called.
let err2 = mgr.auth().await.unwrap_err();
assert!(
matches!(err2, AuthError::Refresh(RefreshTokenError::Permanent(_))),
"second call should return PermanentFailure, got: {err2:?}"
);
// Refresher must have been called exactly once.
assert_eq!(
call_count.load(Ordering::SeqCst),
1,
"refresher must be called exactly once"
);
// hot_swap clears permanent failure; subsequent auth() succeeds.
let valid = GrokAuth {
key: "new-valid-key".into(),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(valid);
assert_eq!(mgr.auth().await.unwrap().key, "new-valid-key");
}
/// auth() re-reads disk via pick_up_sibling_token and returns the
/// sibling-written token when the in-memory token is stale.
#[tokio::test]
async fn auth_legacy_session_picks_up_sibling_disk_token() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let scope = cfg.auth_scope();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
mgr.hot_swap(GrokAuth {
key: "stale-oidc".into(),
auth_mode: AuthMode::Oidc,
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
});
// Sibling writes a valid token to disk.
let fresh = GrokAuth {
key: "fresh-from-sibling".into(),
auth_mode: AuthMode::Oidc,
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
};
let mut store = AuthStore::new();
store.insert(scope, fresh);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
let auth = mgr.auth().await.expect("should pick up sibling token");
assert_eq!(auth.key, "fresh-from-sibling");
}
/// refresh_chain returns TransientFailure when the refresher reports one.
#[tokio::test]
async fn refresh_chain_surfaces_transient_failure() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.hot_swap(GrokAuth {
key: "expired".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
});
struct TransientRefresher;
#[async_trait::async_trait]
impl TokenRefresher for TransientRefresher {
async fn refresh(&self, _: RefreshReason) -> crate::auth::refresh::RefreshOutcome {
crate::auth::refresh::RefreshOutcome::TransientFailure {
message: "idp timeout".into(),
}
}
}
mgr.set_refresher(Arc::new(TransientRefresher));
let err = mgr.auth().await.unwrap_err();
assert!(
matches!(err, AuthError::Refresh(RefreshTokenError::Transient(_))),
"TransientFailure should surface as a transient refresh error, got {err:?}"
);
}
/// Regression: `current()` and `auth()` must agree on whether an
/// expired API key is usable. Pre-fix, `current()` filtered with
/// `!is_token_expired()` (returning None) while the `auth()`
/// `TokenType::ApiKey` branch cloned the stale entry, so the UI saw
/// "logged out" while downstream consumers (trace upload, MCP,
/// embeddings) sent the stale key and hit 401.
#[tokio::test]
async fn auth_returns_expired_api_key_consistently_with_current() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// Seed an API key that is past the 30-day TTL: `create_time` 60
// days ago and no `expires_at`. `is_token_expired` falls through
// to the TTL check and reports `true`.
let expired_key = GrokAuth {
key: "stale-api-key".into(),
auth_mode: AuthMode::ApiKey,
create_time: Utc::now() - Duration::days(60),
expires_at: None,
refresh_token: None,
..GrokAuth::test_default()
};
mgr.hot_swap(expired_key);
// UI / sync read path: the stale key is filtered out.
assert!(
mgr.current().is_none(),
"current() must hide the expired api_key (matches UI/login state)"
);
// Async path: must NOT clone the stale key for downstream
// consumers. Surface `TokenExpiredNoRefresh` so callers can
// funnel the user back through `grok login`.
let err = mgr.auth().await.unwrap_err();
assert!(
matches!(err, AuthError::TokenExpiredNoRefresh),
"auth() must report TokenExpiredNoRefresh for expired api_key, got: {err:?}",
);
assert!(
mgr.get_valid_token().await.is_err(),
"get_valid_token() must error rather than return the stale key"
);
// Sanity: a fresh API key restores both paths.
let fresh_key = GrokAuth {
key: "fresh-api-key".into(),
auth_mode: AuthMode::ApiKey,
create_time: Utc::now(),
expires_at: None,
refresh_token: None,
..GrokAuth::test_default()
};
mgr.hot_swap(fresh_key);
assert_eq!(
mgr.current().map(|a| a.key).as_deref(),
Some("fresh-api-key")
);
assert_eq!(
mgr.get_valid_token().await.ok().as_deref(),
Some("fresh-api-key")
);
}
/// Regression: after a permanent refresh failure (e.g. `invalid_grant`),
/// the proactive refresh task must back off rather than hammer
/// `auth()` in a tight loop. Pre-fix, an expired token + cached
/// PermanentFailure caused `sleep_dur=0` -> `auth()` -> error -> repeat.
///
/// Verified by observing the loop's iteration counter directly: in a
/// 300ms window we tolerate at most a few iterations (one for the
/// initial failure-recording pass, then back-off). Pre-fix the
/// counter would have been in the thousands.
#[tokio::test]
async fn proactive_refresh_backs_off_on_permanent_failure() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// Past-expiry OIDC token: without the backoff guard, the
// proactive loop computes sleep_dur=0 forever.
let expired = GrokAuth {
key: "expired".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(expired);
// Refresher returns invalid_grant the first time it is called and
// counts every invocation. After the first call records the
// permanent failure, the proactive loop must skip subsequent
// calls until the failure is cleared.
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(FailingRefresher {
call_count: call_count.clone(),
}));
let cancel = CancellationToken::new();
mgr.start_proactive_refresh(cancel.clone());
// Give the loop ample time to observe the failure and back off.
// The 300ms window is two orders of magnitude shorter than the
// 5-minute BACKOFF_INTERVAL, so a backed-off loop completes at
// most a couple of iterations: the initial pass that records the
// permanent failure, optionally a few re-check passes if the
// executor races, then sleeps for `BACKOFF_INTERVAL`.
tokio::time::sleep(StdDuration::from_millis(300)).await;
let iterations = mgr.proactive_iteration_count();
let after_failure = call_count.load(Ordering::SeqCst);
// Direct observation of loop progress: a busy-loop produces
// hundreds-to-thousands of iterations in 300ms, the backed-off
// loop produces <= 5.
assert!(
iterations <= 5,
"proactive refresh busy-looped after permanent failure: \
{iterations} iterations (refresher calls: {after_failure})",
);
// Refresher invocation count is a secondary check: the
// permanent_failure short-circuit in `refresh_chain` (added in
// this PR) means at most 1 invocation here.
assert!(
after_failure <= 1,
"refresher must be invoked at most once before the permanent \
failure is recorded, got {after_failure} calls"
);
assert!(
mgr.permanent_failure().is_some(),
"permanent failure must be cached after invalid_grant",
);
// The proactive (background) loop must never emit the manual_auth KPI:
// a background failure is not a user-facing forced re-login.
assert!(
mgr.manual_auth_last_emit().is_none(),
"the proactive background loop must not emit a manual_auth event",
);
cancel.cancel();
}
/// Regression: `start_proactive_refresh` must be
/// idempotent. Calling it twice on the same `Arc<AuthManager>` was
/// previously valid (no guard) and would `tokio::spawn` two
/// background tasks racing on the same in-memory state.
///
/// Asserting on `proactive_iteration_count` is not a meaningful signal
/// because the test fixture (ApiKey + expires_at: None) made every
/// spawned task sleep for `BACKOFF_INTERVAL` immediately. With or
/// without the guard the iteration counter stayed at 0, so that
/// assertion was vacuous (removing the guard left the test passing). The
/// fix is to assert on the new `proactive_start_count()` accessor,
/// which is bumped *inside* the `compare_exchange` success branch
/// in `start_proactive_refresh` -- so it is exactly 1 if the guard
/// fires and N otherwise. This directly observes the invariant
/// instead of inferring it from loop-iteration mechanics.
#[tokio::test]
async fn start_proactive_refresh_is_idempotent() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let stale_api_key = GrokAuth {
key: "stale-api-key".into(),
auth_mode: AuthMode::ApiKey,
create_time: Utc::now() - Duration::days(60),
expires_at: None,
refresh_token: None,
..GrokAuth::test_default()
};
mgr.hot_swap(stale_api_key);
let cancel = CancellationToken::new();
// First call spawns the task; subsequent calls must be no-ops.
mgr.start_proactive_refresh(cancel.clone());
mgr.start_proactive_refresh(cancel.clone());
mgr.start_proactive_refresh(cancel.clone());
// Direct observation of the guard's behavior. Pre-fix: 3.
// Post-fix: exactly 1.
assert_eq!(
mgr.proactive_start_count(),
1,
"start_proactive_refresh idempotency guard failed; expected exactly \
1 spawn after 3 calls",
);
cancel.cancel();
}
/// Proactive path: near-expiry OIDC token -> background task fires
/// refresh_chain(PreRequest) -> consumer sees fresh token.
#[tokio::test]
async fn proactive_refresh_and_consumer_see_fresh_token_end_to_end() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// expires_at inside the 5-min buffer -> proactive fires immediately.
mgr.hot_swap(GrokAuth {
key: "soon-to-expire".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-original".into()),
expires_at: Some(Utc::now() + Duration::seconds(2)),
..GrokAuth::test_default()
});
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: call_count.clone(),
delay: StdDuration::from_millis(0),
}));
let cancel = CancellationToken::new();
mgr.start_proactive_refresh(cancel.clone());
tokio::time::sleep(StdDuration::from_millis(500)).await;
assert!(call_count.load(Ordering::SeqCst) >= 1);
assert_eq!(mgr.get_valid_token().await.unwrap(), "fresh-token");
cancel.cancel();
}
/// Reactive path: expired OIDC token -> try_recover_unauthorized ->
/// refresh_chain(ServerRejected) -> refresher -> consumer sees fresh token.
#[tokio::test]
async fn reactive_401_recovery_produces_fresh_token_end_to_end() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.hot_swap(GrokAuth {
key: "expired-bearer".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-valid".into()),
expires_at: Some(Utc::now() - Duration::minutes(10)),
..GrokAuth::test_default()
});
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: call_count.clone(),
delay: StdDuration::from_millis(0),
}));
assert!(
mgr.try_recover_unauthorized(crate::auth::recovery::RecoverySource::Background)
.await
);
assert_eq!(call_count.load(Ordering::SeqCst), 1);
assert_eq!(mgr.get_valid_token().await.unwrap(), "fresh-token");
}
// refresh_chain permanent-failure short-circuit via recovery is tested
// in recovery::tests::refresh_authority_short_circuits_on_cached_permanent_failure.
/// Different disk RT with expired AT: PermanentFailure is recorded
/// (not demoted to transient), stopping the retry loop.
#[tokio::test]
async fn refresh_chain_records_permanent_failure_when_disk_rt_differs_but_at_expired() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let scope = cfg.auth_scope();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
// Memory has rt-old; disk has rt-new (different RT) but its
// access_token is also expired so try_use_disk_token rejects it
// and we fall through to the refresher.
let stale = GrokAuth {
key: "stale-key".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-old".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
oidc_issuer: Some("https://issuer.example".into()),
oidc_client_id: Some("client-1".into()),
..GrokAuth::test_default()
};
mgr.hot_swap(stale);
let sibling = GrokAuth {
key: "sibling-key".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-new".into()),
expires_at: Some(Utc::now() - Duration::minutes(30)),
oidc_issuer: Some("https://issuer.example".into()),
oidc_client_id: Some("client-1".into()),
..GrokAuth::test_default()
};
let mut store = AuthStore::new();
store.insert(scope, sibling);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
struct FailingRefresher;
#[async_trait::async_trait]
impl crate::auth::refresh::TokenRefresher for FailingRefresher {
async fn refresh(
&self,
_reason: crate::auth::manager::RefreshReason,
) -> crate::auth::refresh::RefreshOutcome {
crate::auth::refresh::RefreshOutcome::permanent(
crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected,
None,
)
}
}
mgr.set_refresher(Arc::new(FailingRefresher));
let err = mgr.auth().await.unwrap_err();
// An expired disk AT means the sibling is dead too — the failure is
// permanent (not demoted to transient). Credentials are retained; the
// scoped verdict is cached and stops the retry storm.
assert!(
matches!(err, AuthError::Refresh(RefreshTokenError::Permanent(_))),
"must surface a permanent failure when disk AT is expired, got: {err:?}",
);
assert!(
mgr.permanent_failure().is_some(),
"verdict must be cached (scoped to the retained credential)",
);
// No-clear invariant: a refresh failure must NOT delete auth.json (a future
// regression that re-adds disk-clear-on-invalid_grant would fail here).
assert!(
mgr.read_disk_auth().is_some(),
"invalid_grant must not delete auth.json (no auto-clear)",
);
// Second attempt short-circuits on the cached verdict — no extra IdP call.
assert!(matches!(
mgr.auth().await.unwrap_err(),
AuthError::Refresh(RefreshTokenError::Permanent(_))
));
}
/// Different disk RT with valid AT: adopt the sibling's token directly.
#[tokio::test]
async fn refresh_chain_demotes_to_transient_when_disk_rt_differs_and_at_valid() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let scope = cfg.auth_scope();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
let stale = GrokAuth {
key: "stale-key".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-old".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
oidc_issuer: Some("https://issuer.example".into()),
oidc_client_id: Some("client-1".into()),
..GrokAuth::test_default()
};
mgr.hot_swap(stale);
let sibling = GrokAuth {
key: "sibling-key".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-new".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
oidc_issuer: Some("https://issuer.example".into()),
oidc_client_id: Some("client-1".into()),
..GrokAuth::test_default()
};
let mut store = AuthStore::new();
store.insert(scope, sibling);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
let calls = Arc::new(AtomicU32::new(0));
struct CountingFailRefresher(Arc<AtomicU32>);
#[async_trait::async_trait]
impl crate::auth::refresh::TokenRefresher for CountingFailRefresher {
async fn refresh(
&self,
_reason: crate::auth::manager::RefreshReason,
) -> crate::auth::refresh::RefreshOutcome {
self.0.fetch_add(1, Ordering::SeqCst);
crate::auth::refresh::RefreshOutcome::permanent(
crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected,
None,
)
}
}
mgr.set_refresher(Arc::new(CountingFailRefresher(calls.clone())));
let result = mgr.auth().await;
assert!(
result.is_ok(),
"should adopt valid sibling token: {result:?}"
);
assert_eq!(result.unwrap().key, "sibling-key");
assert_eq!(
calls.load(Ordering::SeqCst),
0,
"refresher must not be called when disk has a valid token"
);
}
/// Regression: after `clear()` the verdict must *read as absent*
/// — nothing drops it explicitly; it is scoped to the cleared credential and
/// reads through as `None` once that credential is gone — so subsequent
/// `auth()` reports the more useful `NotLoggedIn` (rather than the stale
/// `invalid_grant` from the just-cleared session).
#[tokio::test]
async fn permanent_failure_reads_absent_after_clear_so_auth_reports_not_logged_in() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// Seed + record a permanent failure (as if invalid_grant fired).
let session = GrokAuth {
key: "broken-session".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-revoked".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(session);
record_permanent_failure(
&mgr,
crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected,
);
assert!(mgr.permanent_failure().is_some());
// User runs `grok logout` which calls clear().
mgr.clear().unwrap();
// The diagnostic the user now sees on the next request should be
// "Not logged in. Run `grok login`.", not the stale invalid_grant.
let err = mgr.auth().await.unwrap_err();
assert!(
matches!(err, AuthError::NotLoggedIn),
"auth() after clear() must report NotLoggedIn, got: {err:?}",
);
assert!(
mgr.permanent_failure().is_none(),
"the credential-scoped verdict must read as absent after clear()",
);
// Same check for the hot_swap_clear() path.
let session = GrokAuth {
key: "broken-2".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-2".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(session);
record_permanent_failure(
&mgr,
crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected,
);
mgr.clear_in_memory();
let err = mgr.auth().await.unwrap_err();
assert!(
matches!(err, AuthError::NotLoggedIn),
"auth() after hot_swap_clear() must report NotLoggedIn, got: {err:?}",
);
}
/// `PERMANENT_FAILURE_TTL` means "5 *real* minutes", not "5 awake minutes":
/// a recoverable permanent failure cached just before a system suspend must
/// expire while the machine sleeps. The monotonic clock pauses across suspend,
/// so expiry is judged on both clocks (see `ScopedRefreshFailure::recorded_at`)
/// — this simulates the suspend by rewinding only the wall-clock arm and
/// asserts the failure no longer short-circuits `auth()` on wake.
#[tokio::test]
async fn permanent_failure_expires_on_wall_clock_across_sleep() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// Seed a credential so the verdict scopes to it (an unscoped verdict
// reads through as absent), using the non-sticky `Other` reason — the
// "transient escalation just before lid close" case the TTL exists for.
mgr.hot_swap(GrokAuth {
key: "tok".into(),
..GrokAuth::test_default()
});
record_permanent_failure(&mgr, crate::auth::error::RefreshTokenFailedReason::Other);
assert!(
mgr.permanent_failure().is_some(),
"freshly recorded failure must be live on both clocks",
);
// Simulate a >TTL suspend: monotonic elapsed stays ~0 (paused during
// sleep), wall clock advanced past the TTL.
mgr.force_permanent_failure_wall_aged_out();
assert!(
mgr.permanent_failure().is_none(),
"a slept-through TTL must expire the cached permanent failure on wake",
);
assert!(
!mgr.has_permanent_failure(),
"has_permanent_failure must agree with permanent_failure()",
);
}
// -- Regression: api_key in config.toml must not block OIDC refresh --
/// When a user has an OIDC session (auth.json) AND a model with api_key
/// in config.toml, the OIDC token must still be refreshable. auth()
/// checks TokenType (from AuthManager), not the global auth_method_id.
#[tokio::test]
async fn oidc_refresh_not_blocked_by_model_api_key() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// Expired OIDC token (user has config.toml with api_key on another model).
let expired_oidc = GrokAuth {
key: "expired-session-token".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("valid-rt".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(expired_oidc);
// TokenType is OidcSession regardless of what models exist in config.
assert_eq!(mgr.token_type(), TokenType::OidcSession);
// auth() must attempt OIDC refresh, not short-circuit as ApiKey.
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: call_count.clone(),
delay: StdDuration::from_millis(10),
}));
let result = mgr.auth().await;
assert!(result.is_ok(), "auth() should succeed via OIDC refresh");
assert_eq!(result.unwrap().key, "fresh-token");
assert_eq!(call_count.load(Ordering::SeqCst), 1);
}
// -- direct unit tests for `compute_proactive_sleep` --------
//
// The proactive task's gate chain is a small pure function; testing
// it directly (rather than through `start_proactive_refresh` and a
// sleep window) gives us per-branch coverage that would have caught
// the original vacuity in seconds. Each test below pins one
// arm of `compute_proactive_sleep`.
/// Permanent-failure cached -> backs off (>= BACKOFF_INTERVAL, plus jitter).
#[test]
fn compute_proactive_sleep_permanent_failure_returns_backoff() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let oidc = GrokAuth {
key: "x".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(oidc);
record_permanent_failure(
&mgr,
crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected,
);
let sleep = compute_proactive_sleep(&mgr);
assert!(
sleep >= BACKOFF_INTERVAL && sleep < BACKOFF_INTERVAL + StdDuration::from_secs(60),
"expected backoff + jitter, got {sleep:?}"
);
}
/// Non-refreshable types (LegacySession, ApiKey, None) -> BACKOFF_INTERVAL
/// even when expires_at is past. This is the gate the original
/// test failed to exercise.
#[test]
fn compute_proactive_sleep_non_refreshable_returns_backoff() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// Inject a refresher so the "no refresher" branch doesn't mask
// the gate we're testing.
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: Arc::new(AtomicU32::new(0)),
delay: StdDuration::from_millis(0),
}));
// (a) LegacySession (WebLogin) + Some(past) -- the canonical
// scenario where the absence of the gate produces a busy-loop.
mgr.hot_swap(GrokAuth {
key: "legacy".into(),
auth_mode: AuthMode::WebLogin,
create_time: Utc::now() - Duration::hours(2),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
});
assert_eq!(mgr.token_type(), TokenType::LegacySession);
assert_eq!(compute_proactive_sleep(&mgr), BACKOFF_INTERVAL);
// (b) ApiKey + Some(past).
mgr.hot_swap(GrokAuth {
key: "api".into(),
auth_mode: AuthMode::ApiKey,
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
});
assert_eq!(mgr.token_type(), TokenType::ApiKey);
assert_eq!(compute_proactive_sleep(&mgr), BACKOFF_INTERVAL);
// (c) None (no credentials loaded).
mgr.clear_in_memory();
assert_eq!(mgr.token_type(), TokenType::None);
assert_eq!(compute_proactive_sleep(&mgr), BACKOFF_INTERVAL);
}
/// Sleep gate raised -> BACKOFF_INTERVAL even for a refreshable token past
/// its expiry. Without this gate `refresh_chain` defers every attempt while
/// the proactive loop spins at `sleep_dur=0` (the busy-loop).
#[test]
fn compute_proactive_sleep_sleep_gated_returns_backoff() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: Arc::new(AtomicU32::new(0)),
delay: StdDuration::from_millis(0),
}));
// Refreshable OidcSession past the early-invalidation boundary: without
// the gate this returns 0 (would busy-loop).
mgr.hot_swap(GrokAuth {
key: "oidc".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
});
assert_eq!(
compute_proactive_sleep(&mgr),
StdDuration::from_secs(0),
"precondition: ungated expired refreshable token yields a 0 sleep"
);
mgr.set_system_sleep_imminent(true);
assert_eq!(
compute_proactive_sleep(&mgr),
BACKOFF_INTERVAL,
"sleep gate must back the proactive loop off instead of busy-looping"
);
}
/// Dark wake -> BACKOFF_INTERVAL even for a refreshable token past its expiry.
/// `refresh_chain` defers every attempt during a dark wake (to avoid an IdP
/// refresh straddling an unsignaled re-sleep), so the proactive loop must back
/// off rather than spin at `sleep_dur=0`.
#[test]
fn compute_proactive_sleep_dark_wake_returns_backoff() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: Arc::new(AtomicU32::new(0)),
delay: StdDuration::from_millis(0),
}));
// Refreshable OidcSession past the early-invalidation boundary: without
// the dark-wake gate this returns 0 (would busy-loop).
mgr.hot_swap(GrokAuth {
key: "oidc".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
});
assert_eq!(
compute_proactive_sleep(&mgr),
StdDuration::from_secs(0),
"precondition: non-dark-wake expired refreshable token yields a 0 sleep"
);
mgr.set_dark_wake_for_test(true);
assert_eq!(
compute_proactive_sleep(&mgr),
BACKOFF_INTERVAL,
"dark wake must back the proactive loop off instead of busy-looping"
);
// Returning to a full wake re-enables immediate refresh.
mgr.set_dark_wake_for_test(false);
assert_eq!(
compute_proactive_sleep(&mgr),
StdDuration::from_secs(0),
"full wake must allow the refresh to proceed again"
);
}
/// No refresher configured -> BACKOFF_INTERVAL even for refreshable
/// types. This is the startup-race guard.
#[test]
fn compute_proactive_sleep_no_refresher_returns_backoff() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.hot_swap(GrokAuth {
key: "oidc".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
});
// No `set_refresher` call -- the refresher slot is None.
assert!(mgr.refresher.read().is_none());
assert_eq!(compute_proactive_sleep(&mgr), BACKOFF_INTERVAL);
}
/// Refreshable type + no `expires_at` -> BACKOFF_INTERVAL (the
/// "external binary that doesn't return expiry" case).
#[test]
fn compute_proactive_sleep_refreshable_no_expiry_returns_backoff() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: Arc::new(AtomicU32::new(0)),
delay: StdDuration::from_millis(0),
}));
mgr.hot_swap(GrokAuth {
key: "external".into(),
auth_mode: AuthMode::External,
expires_at: None,
..GrokAuth::test_default()
});
assert_eq!(mgr.token_type(), TokenType::ExternalBinary);
assert_eq!(compute_proactive_sleep(&mgr), BACKOFF_INTERVAL);
}
/// Refreshable type + `Some(past)` and gates pass -> sleep_dur = 0
/// (refresh now). This is the "happy path" the gates don't block.
#[test]
fn compute_proactive_sleep_refreshable_past_expiry_returns_zero() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: Arc::new(AtomicU32::new(0)),
delay: StdDuration::from_millis(0),
}));
mgr.hot_swap(GrokAuth {
key: "oidc".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
});
assert_eq!(mgr.token_type(), TokenType::OidcSession);
assert_eq!(compute_proactive_sleep(&mgr), StdDuration::from_secs(0));
}
/// Refreshable type + `Some(future)` and gates pass -> sleep_dur ~=
/// expires_at - buffer (positive, <= delta). We use a 1-hour horizon
/// and assert the result is in a sane range rather than an exact value
/// (executor scheduling jitter).
#[test]
fn compute_proactive_sleep_refreshable_future_expiry_returns_delta() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: Arc::new(AtomicU32::new(0)),
delay: StdDuration::from_millis(0),
}));
let expires_at = Utc::now() + Duration::hours(1);
mgr.hot_swap(GrokAuth {
key: "oidc".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
expires_at: Some(expires_at),
..GrokAuth::test_default()
});
let dur = compute_proactive_sleep(&mgr);
// Expected: 1h - 5min (early_invalidation) - jitter (060s) ≈ 5455min.
// Range is generous (5159min) to absorb both clock granularity and
// the random jitter added by `compute_proactive_sleep`.
assert!(
dur >= StdDuration::from_secs(51 * 60) && dur <= StdDuration::from_secs(59 * 60),
"expected ~55min, got {dur:?}",
);
}
/// `permanent_failure` cache auto-expires after `PERMANENT_FAILURE_TTL`,
/// so a misclassified transient IdP error (e.g. `invalid_client` during
/// an OAuth client rotation) doesn't permanently log the user out.
#[tokio::test]
async fn permanent_failure_expires_after_ttl() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.hot_swap(GrokAuth {
key: "tok".into(),
..GrokAuth::test_default()
});
record_permanent_failure(
&mgr,
crate::auth::error::RefreshTokenFailedReason::ClientRejected,
);
assert!(
mgr.permanent_failure().is_some(),
"freshly recorded failure should be sticky"
);
mgr.force_permanent_failure_aged_out();
assert!(
mgr.permanent_failure().is_none(),
"aged-out recoverable failure should auto-expire so a retry can succeed"
);
// A revoked refresh token never self-heals: the verdict is sticky past the
// TTL (only a credential change clears it). Stops re-pinging a dead RT.
record_permanent_failure(
&mgr,
crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected,
);
mgr.force_permanent_failure_aged_out();
assert!(
mgr.permanent_failure().is_some(),
"RefreshTokenRejected must stay sticky past the TTL",
);
}
/// The sticky verdict is exempt from BOTH TTL clocks — the monotonic arm
/// (awake time) AND the wall arm (real time across a suspend, added by the
/// sleep-straddle fix). A revoked refresh token never self-heals with time:
/// re-pinging the IdP with it can only fail again, so no amount of aging on
/// either clock may expire the verdict. Only a credential change heals it —
/// the scoped read-through pinned by the `hot_swap` phase below. This is a
/// composition guard: the sticky/non-sticky split and the wall-clock arm
/// landed separately, so neither parent change could test their intersection.
#[tokio::test]
async fn sticky_verdict_survives_both_clocks_but_not_a_credential_change() {
// Guard against a vacuous pass: with < TTL of monotonic uptime the aging
// hook's `checked_sub` no-ops, and a *fresh* verdict would trivially
// satisfy the survival asserts below.
if std::time::Instant::now()
.checked_sub(PERMANENT_FAILURE_TTL + StdDuration::from_secs(1))
.is_none()
{
eprintln!(
"skipping sticky_verdict_survives_both_clocks: host uptime < PERMANENT_FAILURE_TTL"
);
return;
}
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.hot_swap(GrokAuth {
key: "dead".into(),
..GrokAuth::test_default()
});
record_permanent_failure(
&mgr,
crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected,
);
// Age the verdict past the TTL on the monotonic clock AND rewind the
// wall-clock arm past it (what a >TTL suspend looks like to the reader).
mgr.force_permanent_failure_aged_out();
mgr.force_permanent_failure_wall_aged_out();
match mgr.permanent_failure() {
Some(AuthError::Refresh(RefreshTokenError::Permanent(e))) => assert_eq!(
e.reason,
crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected,
"the surviving verdict must carry the sticky reason",
),
other => panic!("sticky verdict must survive both clocks aging out, got {other:?}"),
}
// Time never heals it; a credential change does (read-through, no clear).
mgr.hot_swap(GrokAuth {
key: "fresh".into(),
..GrokAuth::test_default()
});
assert!(
mgr.permanent_failure().is_none(),
"stickiness must not outlive the credential it is scoped to",
);
}
/// The verdict is scoped to the credential that produced it: swapping in a
/// different credential makes it read through as absent, with no explicit
/// clear.
#[tokio::test]
async fn permanent_failure_is_scoped_to_its_credential() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.hot_swap(GrokAuth {
key: "dead".into(),
..GrokAuth::test_default()
});
record_permanent_failure(
&mgr,
crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected,
);
assert!(mgr.permanent_failure().is_some());
// A different credential — no clear call — reads through as no failure.
mgr.hot_swap(GrokAuth {
key: "fresh".into(),
..GrokAuth::test_default()
});
assert!(
mgr.permanent_failure().is_none(),
"verdict must not apply to a different credential",
);
}
/// The verdict is about the *refresh* token: `auth()` must serve a cached
/// access token that is still within its real `expires_at` (buffer-expired
/// but wire-valid) despite a permanent verdict scoped to that credential,
/// without consulting the refresher. Once the same credential passes real
/// expiry, the bypass no longer applies and the permanent error surfaces.
#[tokio::test]
async fn auth_serves_wire_valid_token_despite_permanent_verdict() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// CI runs in K8s pods where is_devbox_environment() is true; without this
// the past-expiry phase would mint via devbox recovery instead of
// surfacing the permanent error.
mgr.set_devbox_env_for_test(false);
// Token in the 5-min buffer (1 min before real expiry): buffer-expired,
// still valid by the IdP's clock.
mgr.hot_swap(GrokAuth {
key: "wire-valid".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-dead".into()),
expires_at: Some(Utc::now() + Duration::minutes(1)),
..GrokAuth::test_default()
});
record_permanent_failure(
&mgr,
crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected,
);
assert!(
mgr.permanent_failure().is_some(),
"verdict must scope to the live credential",
);
// A refresher is wired but must never be consulted: the verdict
// short-circuits the chain and the bypass serves the cached bearer.
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: call_count.clone(),
delay: StdDuration::ZERO,
}));
let served = mgr
.auth()
.await
.expect("a wire-valid token must be served despite the verdict");
assert_eq!(
served.key, "wire-valid",
"auth() must return the cached wire-valid bearer",
);
assert_eq!(
call_count.load(Ordering::SeqCst),
0,
"the verdict must gate the refresher; serving the cached token is free",
);
// Same credential (same key, so the verdict still scopes to it) past its
// real expiry: the bypass no longer applies.
mgr.hot_swap(GrokAuth {
key: "wire-valid".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-dead".into()),
expires_at: Some(Utc::now() - Duration::minutes(1)),
..GrokAuth::test_default()
});
let err = mgr.auth().await.unwrap_err();
assert!(
matches!(err, AuthError::Refresh(RefreshTokenError::Permanent(_))),
"past real expiry the verdict must surface, got: {err:?}",
);
assert_eq!(
call_count.load(Ordering::SeqCst),
0,
"the cached verdict must keep short-circuiting the refresher",
);
}
/// Refresh-failure grace: when the in-memory token is in the 5-min
/// early-invalidation buffer AND `refresh_chain` fails, `auth()`
/// returns the cached token if it's still within its real `expires_at`.
/// The user doesn't see a chat-turn failure for an IdP blip during
/// the buffer window.
#[tokio::test]
async fn auth_returns_cached_token_when_refresh_fails_within_real_expiry() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
// Point at an unreachable proxy so refresh_chain fails fast.
let mgr = Arc::new(AuthManager::new(dir.path(), cfg).with_proxy_base_url("http://127.0.0.1:1"));
// Token in the 5-min buffer (1 min before real expiry) -- past
// the buffer threshold but still valid by the IdP's clock.
let in_buffer = GrokAuth {
key: "still-valid-by-idp".into(),
auth_mode: AuthMode::Oidc,
create_time: Utc::now() - Duration::minutes(55),
user_id: "user-42".into(),
refresh_token: Some("rt".into()),
// Real expiry 1 min away; our 5-min buffer marks it expired.
expires_at: Some(Utc::now() + Duration::minutes(1)),
oidc_issuer: Some("http://127.0.0.1:1".into()),
oidc_client_id: Some("client".into()),
..GrokAuth::test_default()
};
mgr.hot_swap(in_buffer);
let result = mgr.auth().await.expect("grace should return cached token");
assert_eq!(
result.key, "still-valid-by-idp",
"auth() must return the cached token when refresh fails within real expiry"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn update_writes_disk_before_user_enrichment() {
// Mock /user endpoint that blocks on a Notify before responding.
let release = Arc::new(tokio::sync::Notify::new());
let release_for_handler = Arc::clone(&release);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let app = axum::Router::new().route(
"/user",
axum::routing::get(move || {
let r = Arc::clone(&release_for_handler);
async move {
r.notified().await;
axum::Json(serde_json::json!({
"userId": "enriched-user-id",
"email": "enriched@example.com",
"teamId": "enriched-team",
}))
}
}),
);
let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let mgr = Arc::new(
AuthManager::new(dir.path(), cfg).with_proxy_base_url(&format!("http://127.0.0.1:{port}")),
);
// user_id starts empty -- a freshly rotated OIDC token doesn't
// yet know its user_id; that's exactly what /user enriches.
// (If user_id were set AND mismatched the proxy's response, the
// enrichment would correctly bail with reason=user_changed.)
let new_auth = GrokAuth {
key: "rotated-key".into(),
refresh_token: Some("rotated-rt".into()),
user_id: String::new(),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
// `update()` must return well before the `/user` timeout. The
// proxy handler is blocked on `release.notified()` until we say
// so; if `update()` was awaiting `/user` inline, this would
// hang.
let returned = tokio::time::timeout(
std::time::Duration::from_secs(2),
mgr.update(new_auth.clone()),
)
.await
.expect("update() must not block on /user")
.expect("update() must succeed");
assert_eq!(returned.key, "rotated-key");
// Disk must already reflect the rotated tokens, even though
// /user has not responded yet.
let on_disk_before = read_auth_json(&dir.path().join("auth.json")).unwrap();
let entry_before = on_disk_before.values().next().expect("entry written");
assert_eq!(
entry_before.key, "rotated-key",
"rotated key must be on disk before /user lands"
);
assert_eq!(
entry_before.refresh_token.as_deref(),
Some("rotated-rt"),
"rotated refresh_token must be on disk before /user lands"
);
assert_eq!(
entry_before.team_id, None,
"enrichment must not have landed yet"
);
// Now release the /user handler and wait for the enrichment
// task to merge into disk. Poll up to 5s.
release.notify_one();
let auth_path = dir.path().join("auth.json");
let mut enriched = None;
for _ in 0..50 {
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
let store = read_auth_json(&auth_path).unwrap();
let entry = store.values().next().unwrap().clone();
if entry.team_id.is_some() {
enriched = Some(entry);
break;
}
}
let enriched = enriched.expect("enrichment must land within 5s");
// Enrichment must have merged in WITHOUT clobbering the rotated
// tokens.
assert_eq!(enriched.key, "rotated-key", "tokens preserved");
assert_eq!(
enriched.refresh_token.as_deref(),
Some("rotated-rt"),
"refresh_token preserved"
);
assert_eq!(enriched.team_id.as_deref(), Some("enriched-team"));
assert_eq!(enriched.user_id, "enriched-user-id");
server.abort();
}
/// Regression: back-to-back `update()` calls with different
/// `refresh_token`s must converge to the LATEST token on disk, even
/// though both spawned enrichment tasks read-modify-write disk
/// concurrently. This locks the property the spawn-task file lock
/// buys us; without it, the next "drop the lock for performance"
/// PR silently regresses (an interleaved enrichment write can
/// resurrect the older `refresh_token`, re-opening the
/// `invalid_grant` race).
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn enrichment_task_preserves_interleaved_token_rotation() {
// /user returns the SAME user_id for both calls so neither
// enrichment aborts via `user_changed`. The 50 ms latency keeps
// task v1 alive past the v2 update.
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let app = axum::Router::new().route(
"/user",
axum::routing::get(|| async {
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
axum::Json(serde_json::json!({
"userId": "stable-user",
"email": "user@corp.com",
"teamId": "team-alpha",
}))
}),
);
let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let mgr = Arc::new(
AuthManager::new(dir.path(), cfg).with_proxy_base_url(&format!("http://127.0.0.1:{port}")),
);
// Same user_id so neither enrichment aborts; only the rotated
// token fields differ -- the property under test.
let auth_v1 = GrokAuth {
key: "key-v1".into(),
refresh_token: Some("rt-v1".into()),
user_id: "stable-user".into(),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
let auth_v2 = GrokAuth {
key: "key-v2".into(),
refresh_token: Some("rt-v2".into()),
user_id: "stable-user".into(),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
// Two rotations back-to-back. v2's update() lands while v1's
// spawned enrichment task is still in /user.
mgr.update(auth_v1).await.unwrap();
mgr.update(auth_v2).await.unwrap();
// Wait for both spawned tasks to land. Each: 50ms /user + lock
// wait + write. We poll for the eventually-consistent state.
let auth_path = dir.path().join("auth.json");
let mut final_state = None;
for _ in 0..30 {
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
let store = read_auth_json(&auth_path).unwrap();
let entry = store.values().next().unwrap().clone();
// Both rotations done AND enrichment landed.
if entry.refresh_token.as_deref() == Some("rt-v2") && entry.team_id.is_some() {
final_state = Some(entry);
break;
}
}
let final_state = final_state.expect("v2 + enrichment must land within 3s");
// Core invariant: v2's tokens survive both enrichment writes.
assert_eq!(
final_state.refresh_token.as_deref(),
Some("rt-v2"),
"v2 refresh_token must survive v1's stale enrichment write"
);
assert_eq!(
final_state.key, "key-v2",
"v2 access token must survive v1's stale enrichment write"
);
// Enrichment actually ran.
assert_eq!(final_state.team_id.as_deref(), Some("team-alpha"));
assert_eq!(final_state.user_id, "stable-user");
server.abort();
}
/// Regression for the user-switch abort path: if disk's `user_id`
/// changes during an in-flight `/user` call (a different user
/// signed in via a sibling process), the spawned enrichment must
/// abort cleanly rather than overlay a previous user's
/// team/org/profile fields onto the new user's entry.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn enrichment_aborts_when_disk_user_changes_mid_flight() {
// Slow /user so we have time to swap the disk entry mid-flight.
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let app = axum::Router::new().route(
"/user",
axum::routing::get(|| async {
tokio::time::sleep(std::time::Duration::from_millis(300)).await;
axum::Json(serde_json::json!({
"userId": "fetched-user",
"email": "fetched@corp.com",
"teamId": "fetched-team",
}))
}),
);
let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let scope = cfg.auth_scope();
let mgr = Arc::new(
AuthManager::new(dir.path(), cfg.clone())
.with_proxy_base_url(&format!("http://127.0.0.1:{port}")),
);
// Initial entry's user_id matches what /user will return, so
// enrichment WOULD apply normally.
let initial = GrokAuth {
key: "initial-key".into(),
refresh_token: Some("initial-rt".into()),
user_id: "fetched-user".into(),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
mgr.update(initial).await.unwrap();
// Race: while /user is in-flight, a "different user" overwrites
// disk. The enrichment must NOT overlay onto this new entry.
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let intruder = GrokAuth {
key: "intruder-key".into(),
refresh_token: Some("intruder-rt".into()),
user_id: "intruder-user".into(),
team_id: Some("intruder-team".into()),
email: Some("intruder@corp.com".into()),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
let mut store = AuthStore::new();
store.insert(scope.clone(), intruder);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
// The /user mock takes 300 ms; after that the spawned enrichment
// either writes (overlay path -- the regression we're guarding
// against) or aborts silently. Poll the disk over a 3 s window
// and fail fast at the first poll that shows an overlay -- a
// wall-clock `sleep(800ms)` would mask both slow-CI flakes and
// a real regression that just happens to land >800ms in.
let auth_path = dir.path().join("auth.json");
for _ in 0..30 {
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
let store = read_auth_json(&auth_path).unwrap();
let entry = store.get(&scope).expect("entry exists");
assert_eq!(
entry.user_id, "intruder-user",
"intruder's user_id must survive aborted enrichment"
);
assert_eq!(
entry.refresh_token.as_deref(),
Some("intruder-rt"),
"intruder's refresh_token must survive aborted enrichment"
);
assert_eq!(
entry.key, "intruder-key",
"intruder's access token must survive aborted enrichment"
);
assert_eq!(
entry.team_id.as_deref(),
Some("intruder-team"),
"intruder's team must NOT be overwritten with fetched-team"
);
assert_eq!(
entry.email.as_deref(),
Some("intruder@corp.com"),
"intruder's email must NOT be overwritten with fetched@corp.com"
);
}
server.abort();
}
/// Regression: on initial Team-principal login, the OIDC flow
/// stamps `auth.user_id = team_id` as a placeholder so telemetry
/// can distinguish teams immediately (see `extract_user_info` in
/// `oidc.rs`). The `/user` enrichment then returns the *real*
/// user_id and must overlay it onto disk -- this is the entire
/// point of the enrichment call for Team logins. Earlier revisions
/// of this PR compared `disk.user_id` against `user_info.user_id`
/// and treated this legitimate placeholder->real swap as a
/// concurrent user-switch, throwing away the email / team_name /
/// org fields. The guard now compares against the user_id we
/// *wrote* (`auth.user_id`), which matches disk on the bootstrap
/// path and only diverges when a sibling actually stomped.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn enrichment_overlays_team_login_placeholder_user_id() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let app = axum::Router::new().route(
"/user",
axum::routing::get(|| async {
axum::Json(serde_json::json!({
"userId": "real-user-id",
"email": "user@corp.com",
"firstName": "Real",
"lastName": "User",
"principalType": "Team",
"principalId": "team-xyz",
"teamId": "team-xyz",
"teamName": "Some Team",
"teamRole": "MEMBER",
"organizationId": "org-abc",
"organizationName": "Some Org",
"organizationRole": "ORGANIZATION_ROLE_MEMBER",
}))
}),
);
let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let mgr = Arc::new(
AuthManager::new(dir.path(), cfg).with_proxy_base_url(&format!("http://127.0.0.1:{port}")),
);
// Mirrors what `extract_user_info` returns for a Team principal:
// user_id stamped with the team_id placeholder; email + profile
// + team_name + org_* all empty until /user lands.
let team_login = GrokAuth {
key: "team-key".into(),
refresh_token: Some("team-rt".into()),
user_id: "team-xyz".into(),
email: None,
first_name: None,
last_name: None,
principal_type: Some("Team".into()),
principal_id: Some("team-xyz".into()),
team_id: Some("team-xyz".into()),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
mgr.update(team_login).await.unwrap();
// Wait for the spawned enrichment to land.
let auth_path = dir.path().join("auth.json");
let mut enriched = None;
for _ in 0..50 {
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
let store = read_auth_json(&auth_path).unwrap();
let entry = store.values().next().expect("entry exists").clone();
if entry.email.is_some() {
enriched = Some(entry);
break;
}
}
let enriched = enriched.expect("enrichment must overlay onto Team login");
// The whole point: real user_id replaces the team_id placeholder.
assert_eq!(
enriched.user_id, "real-user-id",
"team_id placeholder must be replaced by real user_id from /user"
);
assert_eq!(enriched.email.as_deref(), Some("user@corp.com"));
assert_eq!(enriched.first_name.as_deref(), Some("Real"));
assert_eq!(enriched.last_name.as_deref(), Some("User"));
assert_eq!(enriched.team_name.as_deref(), Some("Some Team"));
assert_eq!(enriched.team_role.as_deref(), Some("MEMBER"));
assert_eq!(enriched.organization_id.as_deref(), Some("org-abc"));
assert_eq!(enriched.organization_name.as_deref(), Some("Some Org"));
assert_eq!(
enriched.organization_role.as_deref(),
Some("ORGANIZATION_ROLE_MEMBER")
);
// Tokens and team-id-as-principal-id preserved.
assert_eq!(enriched.key, "team-key");
assert_eq!(enriched.refresh_token.as_deref(), Some("team-rt"));
assert_eq!(enriched.principal_type.as_deref(), Some("Team"));
assert_eq!(enriched.team_id.as_deref(), Some("team-xyz"));
server.abort();
}
/// Type-system invariant: `apply_user_info_enrichment` must NEVER
/// touch `key`, `refresh_token`, `expires_at`, `oidc_issuer`,
/// `oidc_client_id`, `auth_mode`, `create_time`, or
/// `has_grok_code_access`. The `&mut GrokAuth` signature already
/// enforces this at the type level (you cannot construct a fresh
/// auth from a `UserInfo` -- there's no `From` impl), but a unit
/// test pins the exact list of preserved fields so a future
/// contributor adding a token-like field to both `GrokAuth` and
/// `UserInfo` is forced to look here.
#[test]
fn apply_user_info_enrichment_preserves_token_fields() {
let mut disk = GrokAuth {
key: "ROT_KEY".into(),
refresh_token: Some("ROT_RT".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
oidc_issuer: Some("https://issuer.example".into()),
oidc_client_id: Some("client-xyz".into()),
auth_mode: AuthMode::Oidc,
create_time: Utc::now() - Duration::minutes(10),
has_grok_code_access: Some(true),
user_id: "old-user".into(),
email: Some("old@corp.com".into()),
team_id: Some("old-team".into()),
..GrokAuth::test_default()
};
let snapshot = disk.clone();
let user_info = UserInfo {
user_id: "new-user".into(),
email: Some("new@corp.com".into()),
first_name: Some("New".into()),
last_name: Some("User".into()),
profile_image_asset_id: None,
principal_type: None,
principal_id: None,
team_id: Some("new-team".into()),
team_name: Some("New Team".into()),
team_role: None,
organization_id: None,
organization_name: None,
organization_role: None,
user_blocked_reason: None,
team_blocked_reasons: None,
coding_data_retention_opt_out: None,
subscription_tier: None,
};
apply_user_info_enrichment(&mut disk, user_info);
// Token fields and provenance untouched.
assert_eq!(disk.key, snapshot.key);
assert_eq!(disk.refresh_token, snapshot.refresh_token);
assert_eq!(disk.expires_at, snapshot.expires_at);
assert_eq!(disk.oidc_issuer, snapshot.oidc_issuer);
assert_eq!(disk.oidc_client_id, snapshot.oidc_client_id);
assert_eq!(disk.auth_mode, snapshot.auth_mode);
assert_eq!(disk.create_time, snapshot.create_time);
assert_eq!(disk.has_grok_code_access, snapshot.has_grok_code_access);
// Enrichment fields updated.
assert_eq!(disk.user_id, "new-user");
assert_eq!(disk.email.as_deref(), Some("new@corp.com"));
assert_eq!(disk.team_id.as_deref(), Some("new-team"));
assert_eq!(disk.team_name.as_deref(), Some("New Team"));
assert_eq!(disk.first_name.as_deref(), Some("New"));
}
/// Regression: async provider calls must drive `auth()` so tool requests get refreshed tokens.
#[tokio::test]
#[serial_test::serial] // reaches `resolve_static_api_key`, which reads the key env vars
async fn current_api_key_async_drives_refresh_chain() {
use xai_grok_test_support::EnvGuard;
use xai_grok_tools::types::ApiKeyProvider;
let _xai = EnvGuard::unset("XAI_API_KEY");
let _legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY");
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.hot_swap(GrokAuth {
key: "expired-oidc".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
});
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: call_count.clone(),
delay: StdDuration::from_millis(0),
}));
let provider = super::SharedAuthKeyProvider(mgr.clone());
assert_eq!(provider.current_api_key().as_deref(), Some("expired-oidc"));
let key = provider.current_api_key_async().await;
assert_eq!(key.as_deref(), Some("fresh-token"));
assert_eq!(call_count.load(Ordering::SeqCst), 1);
}
/// Regression: empty or corrupt auth.json must be recoverable on login.
/// Previously the guard in `update()` would skip the disk write on any
/// non-NotFound error, leaving a working in-memory session but a broken file.
#[tokio::test]
async fn update_recovers_from_empty_auth_json() {
let dir = tempfile::tempdir().unwrap();
let auth_path = dir.path().join("auth.json");
let cfg = GrokComConfig::default();
std::fs::write(&auth_path, b"").unwrap();
assert_eq!(std::fs::metadata(&auth_path).unwrap().len(), 0);
let mgr = Arc::new(AuthManager::new(dir.path(), cfg.clone()));
let new_auth = GrokAuth {
key: "recovered-token".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("recovered-rt".into()),
user_id: "recovered-user".into(),
email: Some("user@example.com".into()),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
let result = mgr.update(new_auth.clone()).await;
assert!(
result.is_ok(),
"update must succeed and write to disk: {result:?}"
);
let current = mgr.current();
assert_eq!(
current.as_ref().map(|a| a.key.as_str()),
Some("recovered-token")
);
let on_disk_raw = std::fs::read_to_string(&auth_path).unwrap();
assert!(
!on_disk_raw.is_empty(),
"auth.json must not be empty after recovery"
);
let on_disk: AuthStore =
serde_json::from_str(&on_disk_raw).expect("auth.json must be valid JSON after recovery");
assert!(
on_disk.contains_key(&cfg.auth_scope()),
"persisted scope must be present"
);
assert_eq!(
on_disk.get(&cfg.auth_scope()).map(|a| a.key.as_str()),
Some("recovered-token")
);
}
/// Same as above, but for whitespace-only content.
#[tokio::test]
async fn update_recovers_from_whitespace_only_auth_json() {
let dir = tempfile::tempdir().unwrap();
let auth_path = dir.path().join("auth.json");
let cfg = GrokComConfig::default();
std::fs::write(&auth_path, b" \n\t ").unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg.clone()));
let new_auth = GrokAuth {
key: "ws-token".into(),
auth_mode: AuthMode::Oidc,
user_id: "ws-user".into(),
..make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now())
};
let result = mgr.update(new_auth).await;
assert!(
result.is_ok(),
"update must succeed for whitespace-only file: {result:?}"
);
let on_disk = std::fs::read_to_string(&auth_path).unwrap();
assert!(on_disk.contains("ws-token"), "credential must be persisted");
}
// -- sibling_has_different_refresh_token ----------------------------------
/// Expired disk AT with different RT is not a live sibling.
#[tokio::test]
async fn sibling_different_rt_with_expired_at_is_not_treated_as_live() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg.clone()));
// In-memory: the original RT (revoked via rotation), AT expired.
let original = GrokAuth {
key: "original-at".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-original".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(original);
// Disk: the successor RT from rotation, AT also expired.
let successor = GrokAuth {
key: "successor-at".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-successor".into()),
expires_at: Some(Utc::now() - Duration::minutes(30)),
..GrokAuth::test_default()
};
let mut store = AuthStore::new();
store.insert(cfg.auth_scope(), successor);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
assert!(
!mgr.sibling_has_different_refresh_token(),
"expired disk token must not be treated as a live sibling"
);
}
/// Valid disk AT with different RT is a live sibling.
#[tokio::test]
async fn sibling_different_rt_with_valid_at_is_treated_as_live() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg.clone()));
let original = GrokAuth {
key: "original-at".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-original".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(original);
// Disk: valid token from sibling process.
let sibling = GrokAuth {
key: "sibling-at".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-sibling".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
};
let mut store = AuthStore::new();
store.insert(cfg.auth_scope(), sibling);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
assert!(
mgr.sibling_has_different_refresh_token(),
"valid disk token with different RT must be treated as live sibling"
);
}
/// Regression: refresh_chain(ServerRejected) must bypass the "double-check"
/// early return when the in-memory token is still valid (not expired).
/// Without this, a JWT that is time-valid but missing a subscription claim
/// (post-purchase) is returned as-is and the IdP is never contacted.
#[tokio::test]
async fn refresh_chain_server_rejected_bypasses_valid_token_double_check() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// Seed a valid (non-expired) token — simulates a JWT that is missing
// the subscription claim but is otherwise fine.
let valid_but_rejected = GrokAuth {
key: "pre-subscription-jwt".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-original".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(valid_but_rejected);
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: call_count.clone(),
delay: StdDuration::from_millis(0),
}));
// Confirm the token is considered valid before refresh.
assert_eq!(mgr.current().unwrap().key, "pre-subscription-jwt");
// ServerRejected must force a real refresh despite the token being valid.
let result = mgr
.refresh_chain(
crate::auth::token_type::TokenType::OidcSession,
RefreshReason::ServerRejected,
)
.await;
assert_eq!(
result.unwrap().key,
"fresh-token",
"refresh_chain(ServerRejected) must contact the IdP even with a valid token"
);
assert_eq!(
call_count.load(Ordering::SeqCst),
1,
"refresher must be called exactly once"
);
assert_eq!(
mgr.current().unwrap().key,
"fresh-token",
"in-memory token must be updated to the refreshed one"
);
}
/// When two tasks both get 401 and call refresh_chain(ServerRejected)
/// concurrently, the second caller must return the already-refreshed token
/// without contacting the IdP again. This prevents the double-refresh race
/// where the second caller sends a rotated refresh token → invalid_grant.
#[tokio::test]
async fn refresh_chain_server_rejected_concurrent_skips_redundant_refresh() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// Seed the "rejected" token that both tasks will see.
let rejected = GrokAuth {
key: "rejected-jwt".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-old".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(rejected);
let call_count = Arc::new(AtomicU32::new(0));
// Slow refresher so the second task blocks on the lock long enough
// to observe the first task's refresh result.
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: call_count.clone(),
delay: StdDuration::from_millis(50),
}));
// Both tasks snapshot pre_lock_key = "rejected-jwt", then race for
// the lock. The first refreshes → "fresh-token". The second finds
// current() = "fresh-token" != pre_lock_key → returns early.
let mgr1 = mgr.clone();
let mgr2 = mgr.clone();
let (r1, r2) = tokio::join!(
mgr1.refresh_chain(
crate::auth::token_type::TokenType::OidcSession,
RefreshReason::ServerRejected,
),
mgr2.refresh_chain(
crate::auth::token_type::TokenType::OidcSession,
RefreshReason::ServerRejected,
),
);
// Both must succeed with the refreshed token.
assert_eq!(r1.unwrap().key, "fresh-token");
assert_eq!(r2.unwrap().key, "fresh-token");
// The IdP must be contacted exactly once, not twice.
assert_eq!(
call_count.load(Ordering::SeqCst),
1,
"refresher must be called exactly once; second caller should \
return the already-refreshed token via the double-check guard"
);
}
/// Counterpart: refresh_chain(PreRequest) with a valid token must
/// short-circuit and NOT call the refresher.
#[tokio::test]
async fn refresh_chain_pre_request_short_circuits_on_valid_token() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let valid = GrokAuth {
key: "still-good".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(valid);
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: call_count.clone(),
delay: StdDuration::from_millis(0),
}));
let result = mgr
.refresh_chain(
crate::auth::token_type::TokenType::OidcSession,
RefreshReason::PreRequest,
)
.await;
assert_eq!(result.unwrap().key, "still-good");
assert_eq!(
call_count.load(Ordering::SeqCst),
0,
"PreRequest must NOT call refresher when token is valid"
);
}
// -- login-time inline enrichment -------------------------------------------
/// Axum `/user` stub serving `body`; rejects requests missing `Bearer {token}`.
async fn spawn_user_stub(token: &'static str, body: &'static str) -> String {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let app = axum::Router::new().route(
"/user",
axum::routing::get(move |headers: axum::http::HeaderMap| async move {
let authz = headers
.get("authorization")
.and_then(|v| v.to_str().ok())
.unwrap_or_default();
if authz != format!("Bearer {token}") {
return Err(axum::http::StatusCode::UNAUTHORIZED);
}
Ok(([("content-type", "application/json")], body))
}),
);
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
format!("http://127.0.0.1:{port}")
}
#[tokio::test]
async fn enrich_auth_inline_populates_zdr_flags() {
let body = r#"{"userId":"u-1","teamBlockedReasons":["BLOCKED_REASON_NO_LOGS"],"codingDataRetentionOptOut":true}"#;
let base = spawn_user_stub("tok", body).await;
let dir = tempfile::tempdir().unwrap();
let mgr = AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base);
let mut auth = GrokAuth {
key: "tok".into(),
..GrokAuth::test_default()
};
assert!(!auth.is_data_collection_disabled(), "precondition");
mgr.enrich_auth_inline(&mut auth).await;
assert!(auth.is_zdr_team(), "team_blocked_reasons must be merged");
assert!(auth.coding_data_retention_opt_out);
assert_eq!(auth.user_id, "u-1");
}
#[tokio::test]
async fn enrich_auth_inline_keeps_fields_absent_from_response() {
// `/user` omitting a field must not clear a value the login flow set.
let body = r#"{"userId":"u-1","teamBlockedReasons":["BLOCKED_REASON_NO_LOGS_MODERATED"]}"#;
let base = spawn_user_stub("tok", body).await;
let dir = tempfile::tempdir().unwrap();
let mgr = AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base);
let mut auth = GrokAuth {
key: "tok".into(),
principal_type: Some("Team".into()),
principal_id: Some("team-1".into()),
..GrokAuth::test_default()
};
mgr.enrich_auth_inline(&mut auth).await;
assert_eq!(auth.user_id, "u-1");
assert_eq!(auth.principal_type.as_deref(), Some("Team"));
assert_eq!(auth.principal_id.as_deref(), Some("team-1"));
assert!(auth.is_zdr_team());
assert!(
!auth.coding_data_retention_opt_out,
"absent field stays unchanged"
);
}
#[tokio::test]
async fn enrich_auth_inline_unreachable_server_leaves_auth_unchanged() {
let dir = tempfile::tempdir().unwrap();
// Bind-then-drop to get a port that refuses connections.
let port = {
let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
l.local_addr().unwrap().port()
};
let mgr = AuthManager::new(dir.path(), GrokComConfig::default())
.with_proxy_base_url(&format!("http://127.0.0.1:{port}"));
let mut auth = GrokAuth {
key: "tok".into(),
..GrokAuth::test_default()
};
let before = auth.clone();
mgr.enrich_auth_inline(&mut auth).await;
assert_eq!(auth.user_id, before.user_id);
assert!(!auth.is_data_collection_disabled());
}
// ── force_login_team_uuid spine enforcement ───────────────────────────
//
// Regression coverage for the cached-token bypass: the pin must hold for every
// token the manager hands out (startup, sync reads, `auth()`), not just fresh
// login. Each test fails on the pre-fix tree.
/// `jsonwebtoken` needs a process-level CryptoProvider; tests that encode
/// JWTs can't rely on another test having installed it first.
fn ensure_crypto_provider() {
let _ = jsonwebtoken::crypto::rust_crypto::DEFAULT_PROVIDER.install_default();
}
/// A signed (HS256) access token carrying a `Team` principal, matching the
/// shape `peek_access_token_principal` extracts in production.
fn team_jwt(principal_id: &str) -> String {
ensure_crypto_provider();
jsonwebtoken::encode(
&jsonwebtoken::Header::new(jsonwebtoken::Algorithm::HS256),
&serde_json::json!({
"sub": "user-1",
"principal_type": "Team",
"principal_id": principal_id,
"exp": 9999999999u64,
}),
&jsonwebtoken::EncodingKey::from_secret(b"test-secret"),
)
.unwrap()
}
/// An access token carrying `principal_id` but NO `principal_type`.
fn principal_id_only_jwt(principal_id: &str) -> String {
ensure_crypto_provider();
jsonwebtoken::encode(
&jsonwebtoken::Header::new(jsonwebtoken::Algorithm::HS256),
&serde_json::json!({
"sub": "user-1",
"principal_id": principal_id,
"exp": 9999999999u64,
}),
&jsonwebtoken::EncodingKey::from_secret(b"test-secret"),
)
.unwrap()
}
fn pinned_cfg(team: &str) -> GrokComConfig {
GrokComConfig {
force_login_team_uuid: Some(crate::auth::config::ForceLoginTeam::Single(
team.to_string(),
)),
..GrokComConfig::default()
}
}
/// A valid, non-expired OIDC session whose access token carries `principal_id`.
fn oidc_session_for_team(principal_id: &str) -> GrokAuth {
GrokAuth {
key: team_jwt(principal_id),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
oidc_issuer: Some(crate::auth::config::XAI_OAUTH2_ISSUER.to_string()),
oidc_client_id: Some("client".into()),
..GrokAuth::test_default()
}
}
/// The repro: a wrong-team session persisted to disk (e.g. logged in before
/// the pin was deployed) must be cleared at construction, not silently loaded.
#[test]
fn new_clears_wrong_team_token_loaded_from_disk() {
let dir = tempfile::tempdir().unwrap();
let cfg = pinned_cfg("team-good");
let scope = cfg.auth_scope();
let mut store = AuthStore::new();
store.insert(scope, oidc_session_for_team("team-wrong"));
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
assert!(mgr.current().is_none(), "wrong-team token must be hidden");
assert!(
mgr.current_or_expired().is_none(),
"wrong-team token must be cleared from memory, not just hidden"
);
assert!(
!dir.path().join("auth.json").exists(),
"wrong-team auth.json must be cleared so the next launch re-logs in"
);
}
/// A matching-team session on disk is loaded normally (no false positive).
#[test]
fn new_keeps_matching_team_token_loaded_from_disk() {
let dir = tempfile::tempdir().unwrap();
let cfg = pinned_cfg("team-good");
let scope = cfg.auth_scope();
let tok = oidc_session_for_team("team-good");
let mut store = AuthStore::new();
store.insert(scope, tok.clone());
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
assert_eq!(mgr.current().map(|a| a.key), Some(tok.key));
assert!(dir.path().join("auth.json").exists());
}
/// `auth()` (the wire-bound chokepoint used by pager / MCP /
/// `try_ensure_fresh_auth`) rejects and clears a wrong-team cached token.
#[tokio::test]
async fn auth_rejects_and_clears_wrong_team_cached_token() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), pinned_cfg("team-good")));
// hot_swap bypasses the pin (like a sibling adoption mid-session).
mgr.hot_swap(oidc_session_for_team("team-wrong"));
assert!(mgr.current().is_none(), "sync read must hide the token");
let err = mgr.auth().await.unwrap_err();
assert!(
matches!(err, AuthError::PinnedTeamMismatch { .. }),
"auth() must surface the policy violation, got {err:?}"
);
assert!(
mgr.current_or_expired().is_none(),
"auth() must clear the violating session"
);
}
/// A matching-team cached token flows through `auth()` unchanged.
#[tokio::test]
async fn auth_accepts_matching_team_cached_token() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), pinned_cfg("team-good")));
let tok = oidc_session_for_team("team-good");
mgr.hot_swap(tok.clone());
assert_eq!(mgr.current().map(|a| a.key.clone()), Some(tok.key.clone()));
assert_eq!(mgr.auth().await.unwrap().key, tok.key);
}
/// No pin configured: any team is accepted (the enforcement is opt-in and
/// must not affect default deployments).
#[tokio::test]
async fn no_pin_accepts_any_team_cached_token() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let tok = oidc_session_for_team("team-anything");
mgr.hot_swap(tok.clone());
assert_eq!(mgr.current().map(|a| a.key.clone()), Some(tok.key.clone()));
assert_eq!(mgr.auth().await.unwrap().key, tok.key);
}
/// A token that silently refreshes into a wrong-team principal is rejected by
/// `auth()` (the wrapper gates refresh results, not just the cached fast path).
#[tokio::test]
async fn auth_rejects_token_refreshed_into_wrong_team() {
struct WrongTeamRefresher {
jwt: String,
}
#[async_trait::async_trait]
impl TokenRefresher for WrongTeamRefresher {
async fn refresh(&self, _reason: RefreshReason) -> crate::auth::refresh::RefreshOutcome {
crate::auth::refresh::RefreshOutcome::Success(Box::new(GrokAuth {
key: self.jwt.clone(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-new".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
}))
}
}
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), pinned_cfg("team-good")));
// Expired matching session forces a refresh; the refresher returns a
// wrong-team token (e.g. a re-pinned token family).
mgr.hot_swap(GrokAuth {
expires_at: Some(Utc::now() - Duration::minutes(10)),
..oidc_session_for_team("team-good")
});
mgr.set_refresher(Arc::new(WrongTeamRefresher {
jwt: team_jwt("team-wrong"),
}));
let err = mgr.auth().await.unwrap_err();
assert!(
matches!(err, AuthError::PinnedTeamMismatch { .. }),
"refreshed wrong-team token must be rejected, got {err:?}"
);
}
/// A sibling-written wrong-team token picked up by `force_reload_from_disk`
/// (relay reconnect) is cleared, not just hidden.
#[test]
fn force_reload_clears_wrong_team_token() {
let dir = tempfile::tempdir().unwrap();
let cfg = pinned_cfg("team-good");
let scope = cfg.auth_scope();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg)); // empty disk at startup
let mut store = AuthStore::new();
store.insert(scope, oidc_session_for_team("team-wrong"));
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
mgr.force_reload_from_disk();
assert!(
mgr.current_or_expired().is_none(),
"reloaded wrong-team token must be cleared, not just hidden"
);
assert!(
!dir.path().join("auth.json").exists(),
"force_reload must clear auth.json on a pin violation"
);
}
// -- force_reload_from_disk: transient disk anomaly vs real logout ----------
/// A real incident in miniature: a live in-memory OIDC session (RT
/// present, no permanent_failure) while `auth.json` transiently reads as
/// missing — e.g. the first read right after wake-from-sleep resolves the path
/// to `ENOENT`. The refresh token may exist nowhere else, so the reload must
/// RETAIN it, not discard it (the discard previously kicked off a
/// 401 -> reactive refresh -> suspend-straddle -> invalid_grant cascade).
#[test]
fn force_reload_retains_live_rt_on_transient_file_missing() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let session = GrokAuth {
key: "live-session".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("live-rt".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(session);
assert!(mgr.permanent_failure().is_none());
// No auth.json on disk at all -> FileMissing on every read.
assert!(mgr.read_disk_auth().is_none());
// Zero backoff so the retry budget is exhausted instantly.
mgr.force_reload_from_disk_with(RELOAD_RETRY_TRIES, StdDuration::ZERO);
let retained = mgr.current_or_expired();
assert!(
retained.is_some(),
"a live RT must NOT be discarded on a transient FileMissing",
);
let retained = retained.unwrap();
assert_eq!(retained.key, "live-session");
assert_eq!(retained.refresh_token.as_deref(), Some("live-rt"));
}
/// Contrast with the retain case: once a `permanent_failure` is cached the RT
/// is known-dead, so a persistent FileMissing must drop it (and clear the
/// permanent_failure with it) so the next request reports `NotLoggedIn`.
#[tokio::test]
async fn force_reload_drops_rt_when_permanent_failure_set() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let session = GrokAuth {
key: "broken".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-revoked".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(session);
record_permanent_failure(
&mgr,
crate::auth::error::RefreshTokenFailedReason::RefreshTokenRejected,
);
assert!(mgr.permanent_failure().is_some());
mgr.force_reload_from_disk_with(RELOAD_RETRY_TRIES, StdDuration::ZERO);
assert!(
mgr.current_or_expired().is_none(),
"a known-dead RT (permanent_failure set) must be dropped",
);
assert!(
mgr.permanent_failure().is_none(),
"dropping creds must clear the cached permanent_failure",
);
assert!(matches!(
mgr.auth().await.unwrap_err(),
AuthError::NotLoggedIn
));
}
/// A readable `auth.json` that simply lacks our scope is the trustworthy
/// "logged out / scope removed" signal (distinct from a missing file), so the
/// in-memory credentials are dropped even though an RT is present.
#[test]
fn force_reload_drops_creds_on_entry_missing() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let session = GrokAuth {
key: "live-session".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("live-rt".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(session);
// auth.json exists and is readable, but holds only an unrelated scope ->
// EntryMissing for this manager's scope.
let mut store = AuthStore::new();
store.insert(
"https://example.invalid::nobody".to_string(),
make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now()),
);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
mgr.force_reload_from_disk_with(RELOAD_RETRY_TRIES, StdDuration::ZERO);
assert!(
mgr.current_or_expired().is_none(),
"scope absent on a readable auth.json is a real logout -> drop",
);
}
/// When disk holds a fresh token for our scope, the reload adopts it on the
/// first read (no retry) — the healthy path is unchanged.
#[test]
fn force_reload_adopts_fresh_disk_token() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let scope = cfg.auth_scope();
let mgr = Arc::new(AuthManager::new(dir.path(), cfg));
let expired = GrokAuth {
key: "stale".into(),
refresh_token: Some("old-rt".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
};
mgr.hot_swap(expired);
let fresh = GrokAuth {
key: "fresh-from-disk".into(),
refresh_token: Some("new-rt".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
};
let mut store = AuthStore::new();
store.insert(scope, fresh);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
mgr.force_reload_from_disk_with(RELOAD_RETRY_TRIES, StdDuration::ZERO);
assert_eq!(mgr.current().unwrap().key, "fresh-from-disk");
}
/// A token carrying `principal_id` without `principal_type` is matched on the
/// id alone: the pinned team is accepted, not falsely rejected.
#[tokio::test]
async fn pin_matches_principal_id_without_principal_type() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), pinned_cfg("team-good")));
mgr.hot_swap(GrokAuth {
key: principal_id_only_jwt("team-good"),
auth_mode: AuthMode::Oidc,
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
});
assert!(
mgr.current().is_some(),
"matching team id must be accepted even without principal_type"
);
assert!(mgr.auth().await.is_ok());
}
/// A cached `AuthMode::ApiKey` session is rejected under the kill switch (here
/// implied by a team pin), and honored when it's off.
#[tokio::test]
async fn cached_api_key_session_rejected_when_api_key_auth_disabled() {
let api_key_session = || GrokAuth {
key: "xai-cached-key".into(),
auth_mode: AuthMode::ApiKey,
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
};
// Switch ON (via a team pin, which implies api_key_auth_disabled): reject.
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), pinned_cfg("team-good")));
mgr.hot_swap(api_key_session());
assert!(
mgr.current().is_none(),
"cached api-key session must be hidden under the kill switch"
);
assert!(
matches!(mgr.auth().await, Err(AuthError::ApiKeyAuthDisabled)),
"auth() must reject a cached api-key session under the kill switch"
);
// Switch OFF (no pin / no disable): the api-key session is honored.
let dir2 = tempfile::tempdir().unwrap();
let mgr2 = Arc::new(AuthManager::new(dir2.path(), GrokComConfig::default()));
mgr2.hot_swap(api_key_session());
assert_eq!(
mgr2.current().map(|a| a.key),
Some("xai-cached-key".to_string()),
"api-key session must work normally when the switch is off"
);
}
#[tokio::test]
async fn shared_api_key_provider_resolves_live_bearer() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let auth = GrokAuth {
key: "shared-provider-token".into(),
expires_at: Some(Utc::now() + Duration::hours(1)),
create_time: Utc::now(),
..GrokAuth::test_default()
};
mgr.hot_swap(auth);
let provider = shared_api_key_provider(mgr.clone());
// Synchronous accessor surfaces the current (non-expired) bearer.
assert_eq!(
provider.current_api_key(),
Some("shared-provider-token".to_string()),
"shared_api_key_provider must expose the live bearer to out-of-crate consumers"
);
// Async accessor resolves a valid bearer without a network refresh when
// the cached token is still fresh.
assert_eq!(
provider.current_api_key_async().await,
Some("shared-provider-token".to_string()),
"async accessor must resolve the current bearer for a fresh token"
);
// A hot-swap is reflected on the next resolution (no startup snapshot).
let rotated = GrokAuth {
key: "rotated-token".into(),
expires_at: Some(Utc::now() + Duration::hours(1)),
create_time: Utc::now(),
..GrokAuth::test_default()
};
mgr.hot_swap(rotated);
assert_eq!(
provider.current_api_key(),
Some("rotated-token".to_string()),
"provider must follow the manager's refresh chain rather than snapshot at startup"
);
}
/// No OAuth session → env or auth.json `xai::api_key` for voice/tools.
#[tokio::test]
#[serial_test::serial]
async fn shared_api_key_provider_static_fallthrough() {
use xai_grok_test_support::EnvGuard;
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let provider = shared_api_key_provider(mgr.clone());
{
let _legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY");
let _key = EnvGuard::set("XAI_API_KEY", "env-only-key");
assert_eq!(
provider.current_api_key_async().await.as_deref(),
Some("env-only-key")
);
}
{
let _xai = EnvGuard::unset("XAI_API_KEY");
let _legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY");
crate::auth::store_api_key(dir.path(), "disk-api-key").unwrap();
assert_eq!(
provider.current_api_key_async().await.as_deref(),
Some("disk-api-key")
);
}
{
let _key = EnvGuard::set("XAI_API_KEY", "env-should-lose");
mgr.hot_swap(GrokAuth {
key: "session-bearer".into(),
expires_at: Some(Utc::now() + Duration::hours(1)),
create_time: Utc::now(),
..GrokAuth::test_default()
});
assert_eq!(
provider.current_api_key_async().await.as_deref(),
Some("session-bearer")
);
}
}
#[tokio::test]
#[serial_test::serial]
async fn shared_api_key_provider_kill_switch_blocks_static() {
use xai_grok_test_support::EnvGuard;
let _key = EnvGuard::set("XAI_API_KEY", "blocked");
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(
dir.path(),
GrokComConfig {
disable_api_key_auth: Some(true),
..GrokComConfig::default()
},
));
assert_eq!(
shared_api_key_provider(mgr).current_api_key_async().await,
None
);
}
#[tokio::test]
#[serial_test::serial]
async fn shared_api_key_provider_oidc_preferred_blocks_static() {
use xai_grok_test_support::EnvGuard;
let _key = EnvGuard::set("XAI_API_KEY", "should-not-use");
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(
dir.path(),
GrokComConfig {
preferred_method: Some(crate::auth::PreferredAuthMethod::Oidc),
..GrokComConfig::default()
},
));
assert_eq!(
shared_api_key_provider(mgr).current_api_key_async().await,
None
);
}
/// preferred_method=api_key: leftover session must not beat static API key.
#[tokio::test]
#[serial_test::serial]
async fn shared_api_key_provider_api_key_preferred_skips_session() {
use xai_grok_test_support::EnvGuard;
let _legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY");
let _key = EnvGuard::set("XAI_API_KEY", "static-preferred");
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(
dir.path(),
GrokComConfig {
preferred_method: Some(crate::auth::PreferredAuthMethod::ApiKey),
..GrokComConfig::default()
},
));
mgr.hot_swap(GrokAuth {
key: "leftover-oidc".into(),
expires_at: Some(Utc::now() + Duration::hours(1)),
create_time: Utc::now(),
..GrokAuth::test_default()
});
assert_eq!(
shared_api_key_provider(mgr)
.current_api_key_async()
.await
.as_deref(),
Some("static-preferred")
);
}
/// Expired OAuth must not block static fallthrough on the sync path.
#[tokio::test]
#[serial_test::serial]
async fn shared_api_key_provider_sync_falls_through_when_session_expired() {
use xai_grok_test_support::EnvGuard;
let _legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY");
let _key = EnvGuard::set("XAI_API_KEY", "static-after-expiry");
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.hot_swap(GrokAuth {
key: "expired-oidc".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
});
let provider = shared_api_key_provider(mgr);
assert_eq!(
provider.current_api_key().as_deref(),
Some("static-after-expiry"),
"sync path must not return a dead session token over a live static key"
);
assert_eq!(
provider.current_api_key_async().await.as_deref(),
Some("static-after-expiry")
);
}
/// A session inside the early-invalidation buffer is still wire-valid and
/// must beat a static key on the sync path.
#[tokio::test]
#[serial_test::serial]
async fn shared_api_key_provider_sync_buffered_session_beats_static() {
use xai_grok_test_support::EnvGuard;
use xai_grok_tools::types::ApiKeyProvider;
let _legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY");
let _key = EnvGuard::set("XAI_API_KEY", "leftover-static");
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// Two minutes out: inside the 5-minute buffer, but accepted on the wire.
mgr.hot_swap(GrokAuth {
key: "buffered-oidc".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() + Duration::minutes(2)),
..GrokAuth::test_default()
});
let provider = super::SharedAuthKeyProvider(mgr);
assert_eq!(provider.current_api_key().as_deref(), Some("buffered-oidc"));
}
/// Auth.json create, rewrite (including same-length, caught by the inode in
/// the memo stamp), and logout must all invalidate the disk static-key memo.
#[tokio::test]
#[serial_test::serial]
async fn shared_api_key_provider_disk_memo_follows_rewrites() {
use xai_grok_test_support::EnvGuard;
let _xai = EnvGuard::unset("XAI_API_KEY");
let _legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY");
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let provider = shared_api_key_provider(mgr);
assert_eq!(provider.current_api_key_async().await, None);
for key in ["first-key", "fresh-key", "second-key-rotated"] {
crate::auth::store_api_key(dir.path(), key).unwrap();
assert_eq!(provider.current_api_key_async().await.as_deref(), Some(key));
}
crate::auth::clear_api_key(dir.path()).unwrap();
assert_eq!(provider.current_api_key_async().await, None);
}
Synced from monorepo Changes: - Classify clipboard delivery confidence - Add durable session update append - Scope the xAI session bearer to first-party memory embedding endpoints - Persist subagent outputs to disk and bound long-lived agent state - Add MiniSweAgent:bash for mini-swe-agent parity - Revert taking local sessions off the persistent shell - Contextual tip recommending grok wrap on SSH sessions - Voice STT bearer from model BYOK env_key/api_key - Define exact website policies for sandbox - Gate unsafe shell environments - Shared pin hoist; single require_sha gate for marketplace plugins - Server-signed is-managed claim (closes sidecar-removal downgrade) - Optional require_sha pin for remote plugin installs - Show session title and last exchange in the exit resume hint - Gate shell output redirects - Warn when fail_closed is present but not a boolean - Add canonical text editing core (ratatui-textarea) - Keep execution state out of goal scratch - Add acknowledged persistence primitives - Inherit child network restrictions in sandbox - Fail closed when hook matchers fail to recompile - Add MCP setup preferences for plugin MCPs - Gate sourced shell scripts - Gate file-typed project hooks - grok wrap: restore terminal modes on child death - Harden owner-only permissions on auth and MCP credentials - Create crash dump files with owner-only permissions - Write the agent_id cache owner-only (0600) - SessionMetrics mode skips Mixpanel profile sync - Dashboard: slim live-tail peek - Yank full queued prompt text, not (+N lines) - Defeat clock-rollback on the signed managed-config cache - Stop early session/cancel from overtaking the prompt and wedging the turn slot - Self-heal a diverged agent entrypoint on startup - Add matched inference expectations in test-support - Add AuthSingleFlight cancel/successor gap tests - Remove consumer from external OTEL allowlist and pin scrub coverage - Enable /copy in minimal mode - Surface capacity and API-key detail on 429 errors - Single-flight interactive auth - Fix PageUp/PageDown skipping lines behind sticky prompt header
2026-07-17 14:19:50 +01:00
#[tokio::test]
#[serial_test::serial]
async fn process_key_from_model_env_key() {
use crate::agent::config::{Config, resolve_model_list};
use xai_grok_test_support::EnvGuard;
const ENV: &str = "TEST_MODEL_ENV_KEY";
const TOKEN: &str = "model-env-token";
let _xai = EnvGuard::unset("XAI_API_KEY");
let _legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY");
let _tok = EnvGuard::set(ENV, TOKEN);
let dm = crate::models::default_model();
let cfg = Config::new_from_toml_cfg(
&toml::from_str(&format!(
r#"
[model."{dm}"]
model = "{dm}"
env_key = "{ENV}"
"#
))
.unwrap(),
)
.unwrap();
let key = resolve_model_list(&cfg, None)
.get(dm)
.and_then(|m| m.own_credential())
.unwrap();
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
assert!(mgr.current().is_none());
mgr.set_process_static_api_key(Some(key));
assert_eq!(
shared_api_key_provider(mgr)
.current_api_key_async()
.await
.as_deref(),
Some(TOKEN)
);
}
#[tokio::test]
#[serial_test::serial]
async fn process_key_precedence() {
use xai_grok_test_support::EnvGuard;
let _xai = EnvGuard::unset("XAI_API_KEY");
let _legacy = EnvGuard::unset("GROK_CODE_XAI_API_KEY");
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let provider = shared_api_key_provider(mgr.clone());
assert_eq!(provider.current_api_key_async().await, None);
crate::auth::store_api_key(dir.path(), "disk").unwrap();
assert_eq!(
provider.current_api_key_async().await.as_deref(),
Some("disk")
);
mgr.set_process_static_api_key(Some(" process ".into()));
assert_eq!(
provider.current_api_key_async().await.as_deref(),
Some("process")
);
{
let _key = EnvGuard::set("XAI_API_KEY", "env");
assert_eq!(
provider.current_api_key_async().await.as_deref(),
Some("env")
);
}
mgr.set_process_static_api_key(None);
assert_eq!(
provider.current_api_key_async().await.as_deref(),
Some("disk")
);
let dir_blocked = tempfile::tempdir().unwrap();
let blocked = Arc::new(AuthManager::new(
dir_blocked.path(),
GrokComConfig {
disable_api_key_auth: Some(true),
..GrokComConfig::default()
},
));
blocked.set_process_static_api_key(Some("ignored".into()));
assert_eq!(
shared_api_key_provider(blocked)
.current_api_key_async()
.await,
None
);
}
fn expired_oidc() -> GrokAuth {
GrokAuth {
key: "expired-key".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-old".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
}
}
/// Signals when it has started, then blocks until released.
struct BlockingRefresher {
started: Arc<tokio::sync::Notify>,
release: Arc<tokio::sync::Notify>,
call_count: Arc<AtomicU32>,
}
#[async_trait::async_trait]
impl TokenRefresher for BlockingRefresher {
async fn refresh(&self, _reason: RefreshReason) -> crate::auth::refresh::RefreshOutcome {
self.call_count.fetch_add(1, Ordering::SeqCst);
self.started.notify_one();
self.release.notified().await;
crate::auth::refresh::RefreshOutcome::Success(Box::new(GrokAuth {
key: "fresh-token".into(),
expires_at: Some(Utc::now() + Duration::hours(1)),
refresh_token: Some("rt-new".into()),
..GrokAuth::test_default()
}))
}
}
#[tokio::test]
async fn sleep_gate_defers_refresh_without_calling_idp() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.hot_swap(expired_oidc());
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: call_count.clone(),
delay: StdDuration::from_millis(0),
}));
mgr.set_system_sleep_imminent(true);
let err = mgr.auth().await.unwrap_err();
assert!(
matches!(err, AuthError::Refresh(RefreshTokenError::Transient(_))),
"gated refresh must return a transient refresh error, got {err:?}"
);
assert_eq!(
call_count.load(Ordering::SeqCst),
0,
"the IdP refresher must NOT be called while the sleep gate is raised"
);
}
/// A sleep-deferred refresh must not poison auth state: the deferral is a
/// typed transient (retryable on wake), maps to no `manual_auth` reason (a
/// lid close must never count as a forced re-login in the KPI), and records
/// no permanent-failure verdict — even after more deferred attempts than the
/// refresher-level escalation budget tolerates (the transient-blip budget
/// lives in the refresher, which a deferral never reaches).
///
/// Coverage depth: the gate is raised before the chain starts, so this drives
/// the step-3a deferral. The step-3c pre-IdP re-check (gate raised inside the
/// 3a→3c race window) returns the identical transient error and touches the
/// same state, but is not deterministically reachable without production test
/// hooks, so it is pinned only indirectly by these assertions.
#[tokio::test]
async fn sleep_deferred_refresh_is_transient_no_kpi_no_verdict() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// Pin non-devbox so a deferred refresh surfaces the transient error
// instead of minting via devbox recovery (CI runs in K8s pods).
mgr.set_devbox_env_for_test(false);
mgr.hot_swap(expired_oidc());
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: call_count.clone(),
delay: StdDuration::from_millis(0),
}));
mgr.set_system_sleep_imminent(true);
// More attempts than MAX_CONSECUTIVE_TRANSIENT_FAILURES: deferrals must
// never accrue toward an escalated permanent verdict.
for _ in 0..4 {
let err = mgr.auth().await.unwrap_err();
assert!(
matches!(err, AuthError::Refresh(RefreshTokenError::Transient(_))),
"a sleep-deferred refresh must be transient, got {err:?}"
);
assert_eq!(
crate::auth::recovery::manual_auth_reason(&err),
None,
"a lid-close deferral must never map to a manual_auth KPI reason",
);
}
assert!(
mgr.permanent_failure().is_none(),
"deferrals must not record a permanent-failure verdict",
);
assert_eq!(
call_count.load(Ordering::SeqCst),
0,
"the refresher must never run while the gate is raised",
);
// End-to-end through 401 recovery on a user-facing source: a deferred
// recovery terminates with the transient error and emits no manual_auth.
let mut rec = mgr.unauthorized_recovery(
mgr.current_or_expired(),
crate::auth::recovery::RecoverySource::Turn,
);
let err = rec.next().await.unwrap_err();
assert!(
matches!(err, AuthError::Refresh(RefreshTokenError::Transient(_))),
"deferred recovery must surface the transient deferral, got {err:?}"
);
assert!(
mgr.manual_auth_last_emit().is_none(),
"a sleep-deferred recovery must not emit the manual_auth event",
);
}
#[tokio::test]
async fn dark_wake_defers_refresh_without_calling_idp() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.hot_swap(expired_oidc());
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: call_count.clone(),
delay: StdDuration::from_millis(0),
}));
mgr.set_dark_wake_for_test(true);
let err = mgr.auth().await.unwrap_err();
assert!(
matches!(
err,
AuthError::Refresh(crate::auth::error::RefreshTokenError::Transient(_))
),
"dark-wake refresh must return a transient refresh error, got {err:?}"
);
assert_eq!(
call_count.load(Ordering::SeqCst),
0,
"the IdP refresher must NOT be called during a dark wake (the refresh \
token must not be sent into a possible re-sleep)"
);
// Returning to a full wake lets the refresh proceed and reach the IdP.
mgr.set_dark_wake_for_test(false);
assert_eq!(mgr.auth().await.unwrap().key, "fresh-token");
assert_eq!(
call_count.load(Ordering::SeqCst),
1,
"after a full wake the refresher must be invoked"
);
}
/// A machine stuck reporting a *continuous* dark wake (e.g. an interactive Mac
/// with no display) must not defer refresh forever — once the deferral budget
/// (`DARK_WAKE_DEFER_MAX`) is exhausted, one refresh is forced through. Without
/// this bound the user reaches the same logged-out state the dark-wake guard
/// was added to prevent.
#[tokio::test]
async fn dark_wake_defer_forces_refresh_after_max() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.hot_swap(expired_oidc());
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: call_count.clone(),
delay: StdDuration::from_millis(0),
}));
mgr.set_dark_wake_for_test(true);
// Backdate the start of the deferral run past the bound on both clocks, as
// if we had been continuously in dark wake longer than DARK_WAKE_DEFER_MAX.
let back = super::sleep_gate::DARK_WAKE_DEFER_MAX + StdDuration::from_secs(5);
let (Some(mono), Some(wall)) = (
Instant::now().checked_sub(back),
std::time::SystemTime::now().checked_sub(back),
) else {
return; // machine/clock can't represent the backdate — skip
};
*mgr.dark_wake_defer_since.write() = Some(super::sleep_gate::GateRaise { mono, wall });
assert_eq!(
mgr.auth().await.unwrap().key,
"fresh-token",
"an exhausted dark-wake deferral budget must force the refresh through"
);
assert_eq!(
call_count.load(Ordering::SeqCst),
1,
"the IdP refresher must be invoked once the dark-wake defer budget is exhausted"
);
assert!(
mgr.dark_wake_defer_since.read().is_none(),
"forcing a refresh through must reset the defer budget"
);
}
/// A `DidWake` (`SYSTEM_HAS_POWERED_ON`) event must not reset the dark-wake
/// defer budget while the system is *still* in a dark wake — macOS can deliver
/// powered-on events for dark wakes, and resetting then would stop the budget
/// from ever exhausting, so the forced refresh would never run. Only a genuine
/// full wake clears it.
#[test]
fn dark_wake_defer_budget_survives_powered_on_during_dark_wake() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// Begin a deferral run.
mgr.set_dark_wake_for_test(true);
assert!(
mgr.should_defer_for_dark_wake(),
"a fresh dark wake should defer and start the budget"
);
assert!(mgr.dark_wake_defer_since.read().is_some());
// A powered-on event arrives while still in a dark wake: the budget must
// persist so it can eventually exhaust and force a refresh through.
mgr.set_system_sleep_imminent(false);
assert!(
mgr.dark_wake_defer_since.read().is_some(),
"a powered-on event during a dark wake must not reset the defer budget"
);
// A genuine full wake clears the run.
mgr.set_dark_wake_for_test(false);
mgr.set_system_sleep_imminent(false);
assert!(
mgr.dark_wake_defer_since.read().is_none(),
"a full wake must clear the defer budget"
);
}
/// The `power_listener_started` guard in `is_dark_wake` must short-circuit to
/// `false` when no OS power listener was started (headless / datacenter), so
/// those processes never treat the OS power state as a dark wake. Exercises the
/// guard directly (no dark-wake override installed).
#[test]
fn is_dark_wake_false_when_power_listener_not_started() {
let dir = tempfile::tempdir().unwrap();
let mgr = AuthManager::new(dir.path(), GrokComConfig::default());
assert!(
!mgr.is_dark_wake(),
"is_dark_wake must be false when the power listener was never started"
);
}
#[tokio::test]
async fn sleep_gate_cleared_on_wake_allows_refresh() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.hot_swap(expired_oidc());
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(CountingRefresher {
call_count: call_count.clone(),
delay: StdDuration::from_millis(0),
}));
mgr.set_system_sleep_imminent(true);
mgr.set_system_sleep_imminent(false); // wake
let auth = mgr.auth().await.expect("refresh should succeed after wake");
assert_eq!(auth.key, "fresh-token");
assert_eq!(call_count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn sleep_gate_auto_expires_after_max() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.set_system_sleep_imminent(true);
assert!(mgr.is_sleep_gated(), "freshly-raised gate must be active");
// Simulate a missed wake while awake the whole time: both clocks were
// raised longer ago than the bound.
let back = super::sleep_gate::SLEEP_GATE_MAX + StdDuration::from_secs(5);
let (Some(mono), Some(wall)) = (
Instant::now().checked_sub(back),
std::time::SystemTime::now().checked_sub(back),
) else {
return; // machine/clock can't represent the backdate — not reproducible; skip
};
*mgr.sleep_gate.raised_at.write() = Some(super::sleep_gate::GateRaise { mono, wall });
assert!(
!mgr.is_sleep_gated(),
"a gate older than SLEEP_GATE_MAX must auto-expire"
);
assert!(
mgr.sleep_gate.raised_at.read().is_none(),
"auto-expiry must also lower the gate so a stale state can't linger"
);
}
/// Regression test for the dual-clock backstop: a gate that straddled a real
/// system sleep must auto-expire even though the monotonic clock is still
/// fresh, because the wall clock advanced past the bound during sleep. Before
/// the wall-clock arm this gate stayed shut and an expired token reached the
/// server — the 401 this fix targets.
#[tokio::test]
async fn sleep_gate_auto_expires_when_wall_clock_passes_during_sleep() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.set_system_sleep_imminent(true);
assert!(mgr.is_sleep_gated(), "freshly-raised gate must be active");
// Monotonic clock fresh (as if the machine just slept rather than spending
// the time awake); wall clock pushed past the bound (real time elapsed
// while asleep, where the monotonic clock is frozen).
let back = super::sleep_gate::SLEEP_GATE_MAX + StdDuration::from_secs(5);
let Some(wall) = std::time::SystemTime::now().checked_sub(back) else {
return; // clock can't represent the backdate — not reproducible; skip
};
*mgr.sleep_gate.raised_at.write() = Some(super::sleep_gate::GateRaise {
mono: Instant::now(),
wall,
});
assert!(
!mgr.is_sleep_gated(),
"a gate whose wall-clock age exceeds SLEEP_GATE_MAX must auto-expire \
even though the monotonic clock is still fresh"
);
assert!(
mgr.sleep_gate.raised_at.read().is_none(),
"auto-expiry must also lower the gate so a stale state can't linger"
);
}
#[tokio::test]
async fn sleep_gate_lets_in_flight_refresh_complete() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.hot_swap(expired_oidc());
let started = Arc::new(tokio::sync::Notify::new());
let release = Arc::new(tokio::sync::Notify::new());
let call_count = Arc::new(AtomicU32::new(0));
mgr.set_refresher(Arc::new(BlockingRefresher {
started: started.clone(),
release: release.clone(),
call_count: call_count.clone(),
}));
let m = mgr.clone();
let handle = tokio::spawn(async move { m.auth().await });
started.notified().await;
assert_eq!(
mgr.refresh_in_flight.load(Ordering::SeqCst),
1,
"refresh must be counted as in flight while the IdP call is pending"
);
// `set_system_sleep_imminent` now holds the OS sleep ack until the
// in-flight refresh drains. Drive it from a separate thread — as the real
// OS power-listener thread does — so the tokio runtime stays free to
// complete the refresh while the hold waits.
let sleeper = mgr.clone();
let ack = std::thread::spawn(move || {
let start = Instant::now();
sleeper.set_system_sleep_imminent(true);
start.elapsed()
});
release.notify_one();
let auth = tokio::time::timeout(StdDuration::from_secs(5), handle)
.await
.expect("auth() must return")
.unwrap()
.expect("in-flight refresh must complete, not abort");
let ack_waited = ack.join().expect("ack thread panicked");
assert_eq!(auth.key, "fresh-token");
assert_eq!(call_count.load(Ordering::SeqCst), 1);
assert!(mgr.is_sleep_gated(), "WillSleep must raise the sleep gate");
assert!(
ack_waited < super::sleep_gate::SLEEP_ACK_MAX_WAIT,
"the sleep-ack hold must release when the refresh drains, not wait out \
SLEEP_ACK_MAX_WAIT; waited {ack_waited:?}"
);
assert_eq!(
mgr.refresh_in_flight.load(Ordering::SeqCst),
0,
"in-flight counter must be balanced after completion"
);
}
/// With nothing in flight, the sleep-ack hold must return promptly so the OS
/// suspend is never delayed unnecessarily.
#[test]
fn sleep_ack_hold_returns_immediately_when_nothing_in_flight() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let start = Instant::now();
mgr.test_hold_sleep_ack(StdDuration::from_secs(5));
let waited = start.elapsed();
assert!(
waited < StdDuration::from_millis(250),
"no in-flight refresh must not delay the suspend; waited {waited:?}"
);
}
/// The sleep-ack hold must unblock as soon as the in-flight refresh drains,
/// well before the bound — this is the straddle the fix prevents.
#[test]
fn sleep_ack_hold_releases_when_in_flight_refresh_drains() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.test_enter_refresh_in_flight();
let releaser = mgr.clone();
let drain = std::thread::spawn(move || {
std::thread::sleep(StdDuration::from_millis(120));
releaser.test_exit_refresh_in_flight();
});
let start = Instant::now();
mgr.test_hold_sleep_ack(StdDuration::from_secs(5));
let waited = start.elapsed();
drain.join().unwrap();
assert!(
waited >= StdDuration::from_millis(100),
"must hold the ack until the refresh drains; waited only {waited:?}"
);
assert!(
waited < StdDuration::from_secs(2),
"must release shortly after the drain, not near the bound; waited {waited:?}"
);
assert_eq!(mgr.refresh_in_flight.load(Ordering::SeqCst), 0);
}
/// A refresh that never drains must not pin the machine awake: the hold is
/// bounded and returns at the deadline, leaving the refresh running (never
/// aborted) for the existing straddle telemetry to catch.
#[test]
fn sleep_ack_hold_times_out_when_refresh_never_drains() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
mgr.test_enter_refresh_in_flight(); // never exits
let start = Instant::now();
mgr.test_hold_sleep_ack(StdDuration::from_millis(150));
let waited = start.elapsed();
assert!(
waited >= StdDuration::from_millis(140),
"must wait out the bound; waited only {waited:?}"
);
assert!(
waited < StdDuration::from_secs(1),
"must not exceed the bound by much; waited {waited:?}"
);
assert_eq!(
mgr.refresh_in_flight.load(Ordering::SeqCst),
1,
"the refresh is left running, not aborted, when the hold times out"
);
}
// ── manual_auth KPI ──────────────────────────────────────────
#[test]
fn manual_auth_reason_maps_terminal_and_skips_non_forcing() {
use crate::auth::error::RefreshTokenFailedReason as Reason;
use crate::auth::recovery::manual_auth_reason;
use xai_grok_telemetry::events::ManualAuthReason as R;
let permanent = |reason: Reason| manual_auth_reason(&AuthError::permanent(reason));
// A revoked refresh token forces a re-login -> counts.
assert_eq!(
permanent(Reason::RefreshTokenRejected),
Some(R::RefreshTokenRejected)
);
// Every terminal pipeline error maps to its own bucket (a swapped mapping
// would mis-attribute the KPI).
assert_eq!(
manual_auth_reason(&AuthError::ServerRejectedNoRecovery),
Some(R::NoRefreshAuthority)
);
assert_eq!(
manual_auth_reason(&AuthError::RecoveryExhausted),
Some(R::RecoveryExhausted)
);
assert_eq!(
manual_auth_reason(&AuthError::TokenExpiredNoRefresh),
Some(R::TokenExpiredNoRefresh)
);
assert_eq!(
manual_auth_reason(&AuthError::PinnedTeamMismatch {
message: String::new()
}),
Some(R::WrongTeam)
);
// Self-healing (TTL) reasons, transient / no-credential, and API-key
// lockouts (out of scope for this KPI) don't count.
assert_eq!(permanent(Reason::ClientRejected), None);
assert_eq!(permanent(Reason::Other), None);
assert_eq!(manual_auth_reason(&AuthError::transient("x")), None);
assert_eq!(manual_auth_reason(&AuthError::NotLoggedIn), None);
assert_eq!(manual_auth_reason(&AuthError::ApiKeyAuthDisabled), None);
}
/// Truth table for `relay_should_cancel`: the relay gives up on any terminal
/// auth failure — including `ApiKeyAuthDisabled`, which is deliberately out of
/// the `manual_auth` KPI's scope — and keeps reconnecting through transient
/// blips, absent credentials, and the self-healing permanent reasons (those
/// age out via the TTL, so cancelling on them would orphan a session that
/// recovers minutes later).
#[test]
fn relay_should_cancel_gives_up_only_on_terminal_failures() {
use crate::auth::error::RefreshTokenFailedReason as Reason;
use crate::auth::recovery::relay_should_cancel;
// Terminal: the handshake can't recover; stop reconnecting.
assert!(relay_should_cancel(&AuthError::permanent(
Reason::RefreshTokenRejected
)));
assert!(relay_should_cancel(&AuthError::ServerRejectedNoRecovery));
assert!(relay_should_cancel(&AuthError::RecoveryExhausted));
assert!(relay_should_cancel(&AuthError::TokenExpiredNoRefresh));
assert!(relay_should_cancel(&AuthError::PinnedTeamMismatch {
message: String::new()
}));
// Cancelled even though it never emits the KPI (a kill-switched API key
// means rotate the key, not `/login`).
assert!(relay_should_cancel(&AuthError::ApiKeyAuthDisabled));
// Recoverable: fall through and reconnect.
assert!(!relay_should_cancel(&AuthError::transient("network blip")));
assert!(!relay_should_cancel(&AuthError::permanent(
Reason::ClientRejected
)));
assert!(!relay_should_cancel(&AuthError::permanent(Reason::Other)));
assert!(!relay_should_cancel(&AuthError::NotLoggedIn));
}
// Async so `record` has a runtime for its telemetry `tokio::spawn`: another
// test in the same process can enable the global telemetry client, which would
// otherwise make this emit path panic under a plain `#[test]`.
#[tokio::test]
async fn manual_auth_capture_attributes_and_recorder_debounces() {
use crate::auth::recovery::{ManualAuthTracker, RejectedAuth};
use xai_grok_telemetry::events::{AuthTokenKind, ManualAuthSurface};
let auth = GrokAuth {
key: "dead-token".into(),
user_id: "user-1".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
..GrokAuth::test_default()
};
let snap = RejectedAuth::capture(Some(&auth));
assert_eq!(snap.principal_for_test(), Some("user-1"));
assert_eq!(snap.token_kind_for_test(), AuthTokenKind::OidcSession);
let rec = ManualAuthTracker::default();
let last = || rec.last_token_for_test();
// Records once; a repeat on the same credential debounces.
rec.record(
&snap,
&AuthError::RecoveryExhausted,
ManualAuthSurface::Turn,
);
let id = last();
assert!(id.is_some());
rec.record(
&snap,
&AuthError::PinnedTeamMismatch {
message: String::new(),
},
ManualAuthSurface::Turn,
);
assert_eq!(last(), id);
// A different credential re-arms.
let rearmed = GrokAuth {
key: "another-token".into(),
..auth.clone()
};
let fresh = RejectedAuth::capture(Some(&rearmed));
rec.record(
&fresh,
&AuthError::RecoveryExhausted,
ManualAuthSurface::Turn,
);
assert!(last().is_some() && last() != id);
// A self-healing reason never emits — the KPI counts only forced re-logins.
let healing = ManualAuthTracker::default();
healing.record(
&snap,
&AuthError::permanent(crate::auth::error::RefreshTokenFailedReason::ClientRejected),
ManualAuthSurface::Turn,
);
assert!(healing.last_token_for_test().is_none());
}
/// End-to-end: `next()` emits only for a user-facing, in-scope terminal failure.
/// A credential with no refresh authority terminates with
/// `ServerRejectedNoRecovery` without a refresher.
#[tokio::test]
async fn manual_auth_emits_only_for_user_facing_source() {
use crate::auth::recovery::RecoverySource;
fn mgr_with(dir: &std::path::Path, key: &str, mode: AuthMode) -> Arc<AuthManager> {
let mgr = Arc::new(AuthManager::new(dir, GrokComConfig::default()));
let mut auth = make_auth(Some(Utc::now() + Duration::hours(1)), Utc::now());
auth.user_id = "u1".into();
auth.key = key.into();
auth.auth_mode = mode;
auth.refresh_token = None; // Oidc-sans-refresh-token => LegacySession (in scope)
mgr.hot_swap(auth);
// CI runs in K8s pods where is_devbox_environment() is true; without this
// DevboxRecovery would adopt the seeded valid token and recovery would
// return Ok instead of the terminal ServerRejectedNoRecovery.
mgr.set_devbox_env_for_test(false);
mgr
}
// User-facing + in-scope (legacy session) records.
let d1 = tempfile::tempdir().unwrap();
let turn = mgr_with(d1.path(), "sess-turn", AuthMode::Oidc);
let err = turn
.unauthorized_recovery(turn.current_or_expired(), RecoverySource::Turn)
.next()
.await
.unwrap_err();
assert!(matches!(err, AuthError::ServerRejectedNoRecovery));
// Assert the emitted payload, not just that something fired.
use xai_grok_telemetry::events::{
AuthTokenKind, ManualAuth, ManualAuthReason, ManualAuthSurface,
};
assert_eq!(
turn.manual_auth_last_emit(),
Some(ManualAuth {
reason: ManualAuthReason::NoRefreshAuthority,
trigger: ManualAuthSurface::Turn,
token_kind: AuthTokenKind::LegacySession,
principal: Some("u1".to_string()),
}),
);
// Background source does not record.
let d2 = tempfile::tempdir().unwrap();
let bg = mgr_with(d2.path(), "sess-bg", AuthMode::Oidc);
let _ = bg
.unauthorized_recovery(bg.current_or_expired(), RecoverySource::Background)
.next()
.await;
assert!(bg.manual_auth_last_token().is_none());
// API-key 401 is out of KPI scope even on a user-facing source.
let d3 = tempfile::tempdir().unwrap();
let api = mgr_with(d3.path(), "api-key", AuthMode::ApiKey);
let _ = api
.unauthorized_recovery(api.current_or_expired(), RecoverySource::Turn)
.next()
.await;
assert!(api.manual_auth_last_token().is_none());
}