diff --git a/crates/services/src/auth/cache.rs b/crates/services/src/auth/cache.rs index e097fc129f..a3c10f309d 100644 --- a/crates/services/src/auth/cache.rs +++ b/crates/services/src/auth/cache.rs @@ -13,10 +13,11 @@ const TTL: Duration = Duration::from_secs(30); const MAX_ENTRIES: usize = 1_024; const MAX_BYTES: usize = 4 * 1024 * 1024; const MAX_ENTRY_BYTES: usize = 16 * 1024; -// Remove/insert churn can retain 4096 buckets for 1024 live entries. As in -// the model cache, reserve keys, values and structural overhead for that -// allocation, including storage retained by invalidate(). -const CONTAINER_BYTES: usize = MAX_ENTRIES * 4 * (std::mem::size_of::<([u8; 32], Entry)>() + 64); +// Reserve entry and session-fence buckets, including allocations retained +// after invalidation. Churn can retain four buckets per live entry. +const CONTAINER_BYTES: usize = MAX_ENTRIES + * 4 + * (std::mem::size_of::<([u8; 32], Entry)>() + std::mem::size_of::<([u8; 32], u64)>() + 128); struct Entry { principal: AuthPrincipal, @@ -34,6 +35,8 @@ impl Entry { #[derive(Default)] struct State { generation: u64, + global_invalidation: u64, + session_invalidations: HashMap<[u8; 32], u64>, bytes: usize, entries: HashMap<[u8; 32], Entry>, } @@ -46,6 +49,14 @@ impl AuthenticationCache { Sha256::digest(token.as_bytes()).into() } + fn session_key(user_id: &str, session_id: &str) -> [u8; 32] { + let mut hash = Sha256::new(); + hash.update((user_id.len() as u64).to_be_bytes()); + hash.update(user_id.as_bytes()); + hash.update(session_id.as_bytes()); + hash.finalize().into() + } + pub(super) fn lookup(&self, key: &[u8; 32]) -> (Option, u64) { let mut state = self .0 @@ -80,10 +91,18 @@ impl AuthenticationCache { .unwrap_or_else(|poisoned| poisoned.into_inner()); let now = Instant::now(); let now_unix = chrono::Utc::now().timestamp(); - if state.generation != generation - || now_unix >= expires_at_unix - || now.duration_since(checked_at) >= TTL - { + if state.generation != generation { + let session_invalidated = principal.session_id.as_deref().is_some_and(|session| { + state + .session_invalidations + .get(&Self::session_key(&principal.user.user_id, session)) + .is_some_and(|last| *last > generation) + }); + if generation < state.global_invalidation || session_invalidated { + return false; + } + } + if now_unix >= expires_at_unix || now.duration_since(checked_at) >= TTL { return false; } let bytes = principal_bytes(principal); @@ -125,28 +144,71 @@ impl AuthenticationCache { true } - fn invalidate(&self) { + fn invalidate(&self, session: Option<(&str, &str)>) { let mut state = self .0 .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); state.generation = state.generation.wrapping_add(1); - state.entries.clear(); - state.bytes = 0; + if let Some((user_id, session_id)) = session { + let key = Self::session_key(user_id, session_id); + // Bound the fence metadata. On overflow, fence all older fills + // before discarding their session-specific invalidations. + if !state.session_invalidations.contains_key(&key) + && state.session_invalidations.len() >= MAX_ENTRIES + { + state.global_invalidation = state.generation; + state.session_invalidations.clear(); + } + let generation = state.generation; + state.session_invalidations.insert(key, generation); + state.entries.retain(|_, entry| { + entry.principal.user.user_id != user_id + || entry.principal.session_id.as_deref() != Some(session_id) + }); + state.bytes = state.entries.values().map(|entry| entry.bytes).sum(); + } else { + state.global_invalidation = state.generation; + state.session_invalidations.clear(); + state.entries.clear(); + state.bytes = 0; + } } pub(super) fn mutation(self: &Arc) -> InvalidateOnDrop { - self.invalidate(); - InvalidateOnDrop(Arc::clone(self)) + self.invalidate(None); + InvalidateOnDrop { + cache: Arc::clone(self), + session: None, + } + } + + pub(super) fn session_mutation( + self: &Arc, + user_id: &str, + session_id: &str, + ) -> InvalidateOnDrop { + self.invalidate(Some((user_id, session_id))); + InvalidateOnDrop { + cache: Arc::clone(self), + session: Some((user_id.to_owned(), session_id.to_owned())), + } } } -pub(super) struct InvalidateOnDrop(Arc); +pub(super) struct InvalidateOnDrop { + cache: Arc, + session: Option<(String, String)>, +} impl Drop for InvalidateOnDrop { fn drop(&mut self) { // Includes errors, cancellation, and unknown commit acknowledgements. - self.0.invalidate(); + self.cache.invalidate( + self.session + .as_ref() + .map(|(user, session)| (user.as_str(), session.as_str())), + ); } } @@ -216,6 +278,65 @@ mod tests { assert!(cache.lookup(&key).0.is_none()); } + #[test] + fn session_invalidation_preserves_other_sessions_and_fences_inflight_fills() { + let cache = Arc::new(AuthenticationCache::default()); + let first_key = AuthenticationCache::key("first"); + let second_key = AuthenticationCache::key("second"); + let other_key = AuthenticationCache::key("other"); + let other_fill = AuthenticationCache::key("other-cold-fill"); + let mut first = principal(); + first.session_id = Some("first-session".into()); + let mut second = first.clone(); + second.session_id = Some("second-session".into()); + let mut other = first.clone(); + other.user.user_id = "other-user".into(); + let now = Instant::now(); + let expires = chrono::Utc::now().timestamp() + 60; + for (key, value) in [ + (first_key, &first), + (second_key, &second), + (other_key, &other), + ] { + assert!(cache.publish(key, 0, now, expires, value)); + } + let mutation = cache.session_mutation("owner", "first-session"); + assert!(cache.lookup(&first_key).0.is_none()); + assert_eq!(cache.lookup(&second_key).0, Some(second.clone())); + assert_eq!(cache.lookup(&other_key).0, Some(other.clone())); + assert!(!cache.publish(first_key, 0, now, expires, &first)); + let (_, during) = cache.lookup(&first_key); + drop(mutation); + assert!(!cache.publish(first_key, during, now, expires, &first)); + for _ in 0..100 { + drop(cache.session_mutation("owner", "first-session")); + } + assert!( + cache.publish(other_fill, 0, now, expires, &other), + "a different session's in-flight database observation stays cacheable" + ); + assert_eq!(cache.lookup(&second_key).0, Some(second)); + assert_eq!(cache.lookup(&other_key).0, Some(other)); + } + + #[test] + fn session_fence_overflow_is_bounded_and_fails_closed() { + let cache = Arc::new(AuthenticationCache::default()); + let now = Instant::now(); + let expires = chrono::Utc::now().timestamp() + 60; + for index in 0..=MAX_ENTRIES { + drop(cache.session_mutation("owner", &format!("session-{index}"))); + } + let state = cache.0.lock().unwrap(); + assert!(state.session_invalidations.len() <= MAX_ENTRIES); + assert!(state.global_invalidation > 0); + drop(state); + let key = AuthenticationCache::key("unrelated"); + assert!(!cache.publish(key, 0, now, expires, &principal())); + let (_, current) = cache.lookup(&key); + assert!(cache.publish(key, current, now, expires, &principal())); + } + #[test] fn cache_retention_is_bounded_and_oversized_principals_stay_local() { for content_bytes in [0, MAX_ENTRY_BYTES / 2] { diff --git a/crates/services/src/auth/mod.rs b/crates/services/src/auth/mod.rs index ceb28dc373..2363822843 100644 --- a/crates/services/src/auth/mod.rs +++ b/crates/services/src/auth/mod.rs @@ -1799,33 +1799,26 @@ impl AuthService for DatabaseAuthService { .fetch_one(&mut *tx) .await .map_err(internal_error)?; - // A concurrent disconnect revokes this row under the same user lock; - // the compare-and-revoke prevents refresh from resurrecting it. - let revoked = query( - "UPDATE auth_refresh_tokens SET is_revoked = 1 WHERE token_hash = ? AND is_revoked = 0", + // Replace the current token under the account lock. The old hash no + // longer matches, so refresh and disconnect still have one winner. + let rotated = query( + "UPDATE auth_refresh_tokens SET token_hash = ?, expires_at = ?, created_at = NOW(6), token_prefix = NULL \ + WHERE token_hash = ? AND user_id = ? AND session_id = ? AND is_revoked = 0", ) + .bind(&new_refresh_token_hash) + .bind(expires_at) .bind(&refresh_token_hash) + .bind(&user_id) + .bind(&session_id) .execute(&mut *tx) .await - .map_err(|e| map_auth_sqlx(e, "refresh.revoke_old_token", Some(&pool)))?; - if revoked.rows_affected() != 1 { + .map_err(|e| map_auth_sqlx(e, "refresh.rotate_token", Some(&pool)))?; + if rotated.rows_affected() != 1 { return Err(error_response( StatusCode::UNAUTHORIZED, "Token expired or revoked", )); } - query( - "INSERT INTO auth_refresh_tokens (token_id, user_id, session_id, token_hash, expires_at, is_revoked) \ - VALUES (?, ?, ?, ?, ?, 0)", - ) - .bind(Uuid::new_v4().to_string()) - .bind(&user.user_id) - .bind(&session_id) - .bind(&new_refresh_token_hash) - .bind(expires_at) - .execute(&mut *tx) - .await - .map_err(|e| map_auth_sqlx(e, "refresh.insert_new_refresh_token", Some(&pool)))?; tx.commit() .await .map_err(|e| map_auth_sqlx(e, "refresh.commit_tx", Some(&pool)))?; @@ -1861,42 +1854,50 @@ impl AuthService for DatabaseAuthService { .clone() .ok_or_else(|| error_response(StatusCode::UNAUTHORIZED, "Invalid token"))?; let (session_id, _, _) = self.parse_token_session(&claims)?; + // JWT validation allows clock-skew leeway; logout proofs do not. + if claims.exp <= Utc::now().timestamp() { + return Err(error_response( + StatusCode::UNAUTHORIZED, + "Token expired or revoked", + )); + } let pool = self .get_pool() .await .map_err(|e| map_auth_sqlx(e, "auth.get_pool", None))?; - let refresh_token_hash = sha256_hex(&request.refresh_token); - let stored = self - .fetch_refresh_token(&pool, &refresh_token_hash) + let _invalidation = self + .authentication_cache + .session_mutation(&user_id, &session_id); + let mut tx = pool + .begin() .await - .map_err(|e| map_auth_sqlx(e, "logout.fetch_refresh_token", Some(&pool)))? - .ok_or_else(|| error_response(StatusCode::UNAUTHORIZED, "Token expired or revoked"))?; - - if stored.0 != user_id || stored.2.as_deref() != Some(session_id.as_str()) { + .map_err(|e| map_auth_sqlx(e, "logout.begin_tx", Some(&pool)))?; + let account = query("SELECT user_id FROM auth_users WHERE user_id = ? FOR UPDATE") + .bind(&user_id) + .fetch_optional(&mut *tx) + .await + .map_err(|e| map_auth_sqlx(e, "logout.lock_user", Some(&pool)))?; + if account.is_none() { return Err(error_response( StatusCode::UNAUTHORIZED, - "Token session mismatch", + "Token expired or revoked", )); } - if stored.1 < Utc::now().format("%Y-%m-%dT%H:%M:%S").to_string() { + // Recheck after waiting for the lock. The signed session claim is the + // proof even when a consumed token's old hash has been replaced. + if claims.exp <= Utc::now().timestamp() { return Err(error_response( StatusCode::UNAUTHORIZED, "Token expired or revoked", )); } - - let _invalidation = self.authentication_cache.mutation(); - let mut tx = pool - .begin() - .await - .map_err(|e| map_auth_sqlx(e, "logout.begin_tx", Some(&pool)))?; - query("UPDATE auth_refresh_tokens SET is_revoked = 1 WHERE token_hash = ? AND user_id = ?") - .bind(&refresh_token_hash) + query("UPDATE auth_refresh_tokens SET is_revoked = 1 WHERE user_id = ? AND session_id = ? AND is_revoked = 0") .bind(&user_id) + .bind(&session_id) .execute(&mut *tx) .await - .map_err(|e| map_auth_sqlx(e, "logout.revoke_submitted_token", Some(&pool)))?; + .map_err(|e| map_auth_sqlx(e, "logout.revoke_session", Some(&pool)))?; tx.commit() .await diff --git a/crates/services/tests/auth_cache_db_it.rs b/crates/services/tests/auth_cache_db_it.rs index 1c3f011610..7e38561510 100644 --- a/crates/services/tests/auth_cache_db_it.rs +++ b/crates/services/tests/auth_cache_db_it.rs @@ -29,7 +29,6 @@ struct SnapshotBarrier { impl tracing_subscriber::Layer for Queries { fn on_event(&self, event: &tracing::Event<'_>, _: tracing_subscriber::layer::Context<'_, S>) { if event.metadata().target() == "sqlx::query" { - self.count.fetch_add(1, Ordering::SeqCst); #[derive(Default)] struct Fields(String); impl tracing::field::Visit for Fields { @@ -39,6 +38,11 @@ impl tracing_subscriber::Layer for Queries { } let mut fields = Fields::default(); event.record(&mut fields); + // Pool expansion can initialize a new connection alongside an + // auth query; count application statements, not session setup. + if !fields.0.contains("SET sql_mode=") { + self.count.fetch_add(1, Ordering::SeqCst); + } if fields.0.contains("SELECT u.user_id, u.username") && self.snapshot.armed.swap(false, Ordering::SeqCst) { @@ -175,6 +179,26 @@ async fn jwt_authentication_uses_one_cold_query_and_zero_warm_queries_with_local .current_principal(&second) .await .expect("logout does not revoke another device"); + assert_eq!( + count.swap(0, Ordering::SeqCst), + 0, + "another session stays warm" + ); + for _ in 0..3 { + service + .logout(AuthRefreshRequestData { + refresh_token: tokens[0].refresh_token.clone(), + }) + .await + .unwrap(); + count.store(0, Ordering::SeqCst); + service.current_principal(&second).await.unwrap(); + assert_eq!( + count.swap(0, Ordering::SeqCst), + 0, + "replayed logout must not flush unrelated sessions" + ); + } sqlx::query("UPDATE auth_users SET is_active = 0 WHERE user_id = ?") .bind(&user.user_id) @@ -246,8 +270,8 @@ async fn jwt_authentication_uses_one_cold_query_and_zero_warm_queries_with_local let subsequent = observer.current_principal(&second).await; assert_eq!( count.swap(0, Ordering::SeqCst), - 1, - "invalidated fill cannot re-enter cache" + usize::from(same_session), + "only the revoked session's in-flight fill must be fenced" ); if same_session { assert_eq!( diff --git a/crates/services/tests/auth_logout_rotation_db_it.rs b/crates/services/tests/auth_logout_rotation_db_it.rs new file mode 100644 index 0000000000..0040d0503f --- /dev/null +++ b/crates/services/tests/auth_logout_rotation_db_it.rs @@ -0,0 +1,369 @@ +mod common; + +use astra_core::JwtSettings; +use astra_services::{ + AuthLoginRequestData, AuthRefreshRequestData, AuthRegisterRequestData, AuthService, + DatabaseAuthService, RetentionPolicy, cleanup_expired_data, +}; +use axum::http::{HeaderMap, HeaderValue, StatusCode}; +use std::sync::{ + Arc, Condvar, Mutex, + atomic::{AtomicBool, Ordering}, +}; +use tracing_subscriber::prelude::*; + +#[derive(Default)] +struct AccountLockGate { + armed: AtomicBool, + acquired: tokio::sync::Notify, + released: Mutex, + wake: Condvar, +} + +impl AccountLockGate { + fn arm(&self) { + *self.released.lock().unwrap() = false; + self.armed.store(true, Ordering::SeqCst); + } + + fn release(&self) { + *self.released.lock().unwrap() = true; + self.wake.notify_all(); + } +} + +struct PauseAccountLock(Arc); + +impl tracing_subscriber::Layer for PauseAccountLock { + fn on_event(&self, event: &tracing::Event<'_>, _: tracing_subscriber::layer::Context<'_, S>) { + if event.metadata().target() != "sqlx::query" || !self.0.armed.load(Ordering::SeqCst) { + return; + } + #[derive(Default)] + struct Fields(String); + impl tracing::field::Visit for Fields { + fn record_debug(&mut self, _: &tracing::field::Field, value: &dyn std::fmt::Debug) { + self.0.push_str(&format!("{value:?}")); + } + } + let mut fields = Fields::default(); + event.record(&mut fields); + if fields + .0 + .contains("SELECT user_id FROM auth_users WHERE user_id = ? FOR UPDATE") + && self.0.armed.swap(false, Ordering::SeqCst) + { + self.0.acquired.notify_one(); + tokio::task::block_in_place(|| { + let (released, timeout) = self + .0 + .wake + .wait_timeout_while( + self.0.released.lock().unwrap(), + std::time::Duration::from_secs(20), + |released| !*released, + ) + .unwrap(); + assert!( + *released && !timeout.timed_out(), + "account lock barrier timed out" + ); + }); + } + } +} + +fn headers(token: &str) -> HeaderMap { + let mut headers = HeaderMap::new(); + headers.insert( + "authorization", + HeaderValue::from_str(&format!("Bearer {token}")).unwrap(), + ); + headers +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "requires live MatrixOne: ASTRA_TEST_DB_IT=1"] +async fn logout_revokes_the_session_across_refresh_rotation_without_revoking_other_sessions() { + let gate = Arc::new(AccountLockGate::default()); + tracing::subscriber::set_global_default( + tracing_subscriber::registry().with(PauseAccountLock(gate.clone())), + ) + .expect("this test binary owns its tracing subscriber"); + let (pool, settings) = common::setup_pool_and_settings().await; + let service = DatabaseAuthService::new( + settings, + JwtSettings { + secret_key: "isolated-logout-rotation-test-key".into(), + algorithm: "HS256".into(), + access_token_expire_minutes: 5, + refresh_token_expire_days: 1, + }, + ) + .with_pool(pool.clone()); + let username = format!("logout_{}", uuid::Uuid::new_v4().simple()); + let password = "correct-horse-battery-staple"; + let user = service + .register(AuthRegisterRequestData { + username: username.clone(), + email: format!("{username}@example.test"), + password: password.into(), + display_name: None, + }) + .await + .unwrap(); + let login = || AuthLoginRequestData { + username: username.clone(), + password: password.into(), + }; + let original = service.login(login()).await.unwrap(); + // Model a still-valid token issued more than the cleanup retention age ago. + sqlx::query("UPDATE auth_refresh_tokens SET created_at = DATE_SUB(NOW(6), INTERVAL 8 DAY) WHERE user_id = ?") + .bind(&user.user_id) + .execute(pool.get()) + .await + .unwrap(); + let other = service.login(login()).await.unwrap(); + let key = b"isolated-logout-rotation-test-key"; + let mut expired_claims = jsonwebtoken::decode::( + &original.refresh_token, + &jsonwebtoken::DecodingKey::from_secret(key), + &jsonwebtoken::Validation::default(), + ) + .unwrap() + .claims; + // This is inside the JWT library's usual expiry leeway. + expired_claims["exp"] = (chrono::Utc::now().timestamp() - 1).into(); + let expired = jsonwebtoken::encode( + &jsonwebtoken::Header::default(), + &expired_claims, + &jsonwebtoken::EncodingKey::from_secret(key), + ) + .unwrap(); + for invalid in [ + original.access_token.clone(), + "invalid.token".to_string(), + expired, + ] { + assert_eq!( + service + .logout(AuthRefreshRequestData { + refresh_token: invalid, + }) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + } + let mut rotated = service + .refresh(AuthRefreshRequestData { + refresh_token: original.refresh_token.clone(), + }) + .await + .unwrap(); + assert_eq!( + service + .refresh(AuthRefreshRequestData { + refresh_token: original.refresh_token.clone() + }) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + // Repeated rotations replace the same session row rather than retaining a + // historical row for every consumed token. + for _ in 0..3 { + rotated = service + .refresh(AuthRefreshRequestData { + refresh_token: rotated.refresh_token, + }) + .await + .unwrap(); + } + let cleanup = cleanup_expired_data(pool.get(), &RetentionPolicy::default()) + .await + .unwrap(); + assert!( + cleanup + .iter() + .any(|result| result.table == "auth_refresh_tokens") + ); + let retained: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM auth_refresh_tokens WHERE user_id = ?") + .bind(&user.user_id) + .fetch_one(pool.get()) + .await + .unwrap(); + assert_eq!(retained, 2, "two sessions retain only one refresh row each"); + service + .logout(AuthRefreshRequestData { + refresh_token: original.refresh_token.clone(), + }) + .await + .unwrap(); + // An old signed proof still revokes the current row even though its own + // hash disappeared on rotation. Revoked rows are eligible after retention. + sqlx::query("UPDATE auth_refresh_tokens SET created_at = DATE_SUB(NOW(6), INTERVAL 8 DAY) WHERE user_id = ? AND is_revoked = 1") + .bind(&user.user_id) + .execute(pool.get()) + .await + .unwrap(); + let cleanup = cleanup_expired_data(pool.get(), &RetentionPolicy::default()) + .await + .unwrap(); + assert!( + cleanup + .iter() + .any(|result| result.table == "auth_refresh_tokens" && result.rows_deleted > 0) + ); + assert_eq!( + service + .refresh(AuthRefreshRequestData { + refresh_token: rotated.refresh_token + }) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + assert_eq!( + service + .current_principal(&headers(&rotated.access_token)) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + service + .current_principal(&headers(&other.access_token)) + .await + .expect("another session remains active"); + + let concurrent = service.login(login()).await.unwrap(); + let (left, right) = tokio::join!( + service.refresh(AuthRefreshRequestData { + refresh_token: concurrent.refresh_token.clone() + }), + service.refresh(AuthRefreshRequestData { + refresh_token: concurrent.refresh_token.clone() + }), + ); + let next = match (left, right) { + (Ok(next), Err(error)) | (Err(error), Ok(next)) => { + assert_eq!(error.0, StatusCode::UNAUTHORIZED); + next + } + _ => panic!("a single-use refresh token must have exactly one winner"), + }; + service + .logout(AuthRefreshRequestData { + refresh_token: concurrent.refresh_token, + }) + .await + .unwrap(); + assert_eq!( + service + .refresh(AuthRefreshRequestData { + refresh_token: next.refresh_token + }) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + + // Pause the winner after it acquires the real account lock so both orders + // are exercised without depending on task scheduling. + for refresh_first in [true, false] { + let token = service.login(login()).await.unwrap(); + gate.arm(); + let winner_service = service.clone(); + let winner_token = token.refresh_token.clone(); + let winner = tokio::spawn(async move { + if refresh_first { + winner_service + .refresh(AuthRefreshRequestData { + refresh_token: winner_token, + }) + .await + .map(Some) + } else { + winner_service + .logout(AuthRefreshRequestData { + refresh_token: winner_token, + }) + .await + .map(|()| None) + } + }); + tokio::time::timeout(std::time::Duration::from_secs(15), gate.acquired.notified()) + .await + .unwrap(); + let other_service = service.clone(); + let other_token = token.refresh_token.clone(); + let mut loser = tokio::spawn(async move { + if refresh_first { + other_service + .logout(AuthRefreshRequestData { + refresh_token: other_token, + }) + .await + .map(|()| None) + } else { + other_service + .refresh(AuthRefreshRequestData { + refresh_token: other_token, + }) + .await + .map(Some) + } + }); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), &mut loser) + .await + .is_err() + ); + gate.release(); + let first = winner.await.unwrap().unwrap(); + let second = loser.await.unwrap(); + if let Some(next) = first { + second.unwrap(); + assert_eq!( + service + .refresh(AuthRefreshRequestData { + refresh_token: next.refresh_token + }) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + assert_eq!( + service + .current_principal(&headers(&next.access_token)) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + } else { + assert_eq!(second.unwrap_err().0, StatusCode::UNAUTHORIZED); + } + service + .current_principal(&headers(&other.access_token)) + .await + .expect("another session remains active"); + } + for table in ["auth_refresh_tokens", "auth_user_roles", "auth_users"] { + sqlx::query(sqlx::AssertSqlSafe(format!( + "DELETE FROM {table} WHERE user_id = ?" + ))) + .bind(&user.user_id) + .execute(pool.get()) + .await + .unwrap(); + } + pool.get().close().await; +} diff --git a/docs/design/authentication.md b/docs/design/authentication.md index f50b9442f4..5c17b9f988 100644 --- a/docs/design/authentication.md +++ b/docs/design/authentication.md @@ -55,9 +55,11 @@ is retained in the cache. The cache is shared by service clones and bounded to 1,024 entries, 4 MiB of accounted storage and 16 KiB per entry. Reconfiguring a service's pool, issuer or -encryptor starts a new namespace. Logout, disconnect and Memoria binding changes -invalidate local observations and fence in-flight fills, including uncertain -commits. Other processes and out-of-band account/session mutations can take up +encryptor starts a new namespace. Logout invalidates only its session's local +observations and fences matching in-flight fills; account-wide changes still +invalidate all observations, including on uncertain commits. Session fences are +bounded within the cache; overflow conservatively fences older in-flight fills. +Other processes and out-of-band account/session mutations can take up to 30 seconds to affect ordinary request admission. Expired observations require a successful fresh lookup; database errors do not enable stale fallback. @@ -67,6 +69,26 @@ unrelated principal. Absolute token, observed Session and Memoria access expiry are checked again after awaited validation. A request starting after a local logout cannot reuse an invalidated observation. +Refresh and logout serialize on the account lock. Refresh atomically replaces +the active token hash in its existing session row; the consumed hash cannot +refresh again. Logout uses a signed, strictly unexpired refresh token's user and +session claims to revoke that session, even after the old hash is replaced or +cleaned up. A valid proof with no remaining session is an idempotent success; +it cannot revoke another session. This relies on unique session IDs and trusted +JWT signing keys: removing an old hash alone no longer revokes its logout proof. +Expired or revoked rows remain eligible for the existing bounded cleanup. +Do not mix this version with older servers: an old server may report logout +success after a new server has already replaced the hash, leaving the session +active. For upgrade or rollback, stop admission to all API replicas, drain or +terminate their in-flight requests, and wait until every old process has exited +before starting any replica of the other version. This cutover interrupts +service; ordinary rolling updates and session affinity do not provide the +required ordering. Keep all API replicas on the same side of the cutover. +After rollback, an older server cannot use a token already consumed by the +newer version to revoke its successor; it rejects that logout. Draining +requests does not restore replaced hashes, so clients must use their current +refresh tokens or the affected sessions must be revoked separately. + This cache observes identity, not operation permissions. Memory operations continue resolving scoped credentials through the canonical credential resolver; provider-request replay checks, sensitive reauthentication and execution-handoff diff --git a/docs/reference/api-reference.md b/docs/reference/api-reference.md index 527aac18fe..368db7659a 100644 --- a/docs/reference/api-reference.md +++ b/docs/reference/api-reference.md @@ -10,6 +10,11 @@ All protected endpoints require a JWT. The public authentication exceptions are `/live`, `/ready`, `/health`, `/auth/register`, `/auth/login`, `GET /auth/methods`, `POST /auth/memoria`, `/auth/refresh`, and `/auth/logout` (refresh/logout authenticate the supplied refresh token): +Logout revokes the token's session, including a successor issued by a +concurrent refresh. A signed, unexpired refresh token can terminate its own +session even after rotation; it cannot be used to refresh again. A valid proof +for a session already gone is an idempotent success. Other sessions are unaffected. + ``` Authorization: Bearer ``` diff --git a/packages/sdk/README.md b/packages/sdk/README.md index 16be074fbc..d385bf1160 100644 --- a/packages/sdk/README.md +++ b/packages/sdk/README.md @@ -161,6 +161,53 @@ already accepted frame to other subscribed callbacks. ### Browser apps and the BFF pattern +Within one `AstraClient`, concurrent REST, live streams and buffered run-event requests share +one token refresh. A delayed 401 reuses the rotated token. Refresh responses +started before `setTokens`, successful login/registration, or logout cannot +restore the previous credentials or retry an old request under a new account. +`setTokens(accessToken)` clears any previous refresh token; pass both tokens +when replacing a credential pair. Pending login and logout operations cannot +overwrite a newer authentication decision. +This coordination is per client instance, not across separate BFF requests. + +`onTokenRefresh` may be asynchronous and call the same client. Requests already +awaiting its rotation still wait for the callback, so the callback must not +await a request that has already joined that rotation. A delayed old-token 401 +returns its original error while that notification is pending or has failed +without a recovery in progress, so it cannot use credentials whose persistence +is uncertain. Callback failure +preserves the rotated credentials in memory. Requests made during notification +can use the current token, but a 401 does not start another refresh until the +notification finishes. This prevents overlapping persistence callbacks from +storing a superseded refresh token. A delayed 401 joins a newer in-flight +rotation when no notification is active. +If credentials can be replaced while a callback writes to external storage, +that storage must use an atomic conditional write or another ordering mechanism +to prevent an older callback from overwriting newer credentials. Checking +`getTokens()` before an asynchronous write alone is not sufficient. +The Web runtime client persists a captured pair only if it still matches that +instance's current credentials after asynchronous cookie acquisition. + +The Web runtime's raw HTTP requests share this SDK refresh owner. Its API URL +is origin-relative: a pathname suffix is not a gateway prefix. Standalone SDK +`baseUrl` and `pathPrefix` routing are unchanged. Web refresh or notification +failure returns the original 401, which its typed adapter converts to a runtime +error; initial transport errors still propagate. + +`fetchResponse(path, init, { auth })` returns an unconsumed response, including +non-2xx responses, without inferring a JSON content type. Its default auth is +`"optional"`; `"none"` suppresses default authorization and automatic refresh. +Explicit request headers still override defaults and are not automatically +refreshed using the client's credentials. Raw mutation requests are not retried +unless `retryUnauthorized` is explicitly enabled for an endpoint that rejects +authentication before admitting the operation. Streams and one-shot bodies +are never replayed after acceptance. `getTokens()` returns a +detached snapshot of current credentials, not the initial configuration. +Bodies must remain replayable for the single automatic retry. + +Buffered `getRunEvents` uses the same authenticated request and structured +`AstraApiError` decoding as REST, including error codes and action hints. + Do not expose long-lived **refresh tokens** in browser-accessible JavaScript. The dashboard (`web/`) keeps tokens in **httpOnly cookies** and routes browser requests through explicit **Next.js Route Handlers** such as `/api/chats`, diff --git a/packages/sdk/src/__tests__/auth-lifecycle-regression.test.ts b/packages/sdk/src/__tests__/auth-lifecycle-regression.test.ts new file mode 100644 index 0000000000..14ba17e8cd --- /dev/null +++ b/packages/sdk/src/__tests__/auth-lifecycle-regression.test.ts @@ -0,0 +1,235 @@ +import { AstraClient } from '../client'; + +function deferred() { + let resolve!: (value: T) => void; + const promise = new Promise((done) => { resolve = done; }); + return { promise, resolve }; +} +const unauthorized = () => Response.json({ detail: 'expired' }, { status: 401 }); +const tokens = { access_token: 'new-access', refresh_token: 'new-refresh' }; +const bearer = (init?: RequestInit) => new Headers(init?.headers).get('authorization'); +afterEach(() => vi.unstubAllGlobals()); + +test('replacing only the access token cannot refresh as the previous account', async () => { + const fetch = vi.fn(async () => unauthorized()); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'alice', refreshToken: 'alice-refresh' }); + client.setTokens('bob'); + await expect(client.getMe()).rejects.toMatchObject({ status: 401 }); + expect(client.getTokens()).toEqual({ accessToken: 'bob', refreshToken: null }); + expect(fetch).toHaveBeenCalledTimes(1); +}); + +test.each([204, 500])('late logout response %i cannot erase a newer login', async (status) => { + const logout = deferred(); + vi.stubGlobal('fetch', vi.fn(async (url: string) => { + if (url.endsWith('/auth/logout')) return logout.promise; + if (url.endsWith('/auth/login')) return Response.json(tokens); + throw new Error(url); + })); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const pending = client.logout().then(() => 'ok', () => 'error'); + expect(client.getTokens()).toEqual({ accessToken: null, refreshToken: null }); + await client.login('bob', 'password'); + logout.resolve(status === 204 ? new Response(null, { status }) : new Response('failed', { status })); + expect(await pending).toBe(status === 204 ? 'ok' : 'error'); + expect(client.getTokens()).toEqual({ accessToken: 'new-access', refreshToken: 'new-refresh' }); +}); + +test.each(['login', 'register'] as const)('pending %s cannot restore a session after logout', async (method) => { + const result = deferred(); + vi.stubGlobal('fetch', vi.fn(async () => result.promise)); + const client = new AstraClient({ baseUrl: 'https://example.test' }); + const pending = client[method]('bob', 'password'); + await client.logout(); + result.resolve(Response.json(tokens)); + await expect(pending).rejects.toThrow('superseded'); + expect(client.getTokens()).toEqual({ accessToken: null, refreshToken: null }); +}); + +test('a slower login cannot overwrite a later successful login', async () => { + const first = deferred(); + let calls = 0; + vi.stubGlobal('fetch', vi.fn(async () => ++calls === 1 ? first.promise : Response.json(tokens))); + const client = new AstraClient({ baseUrl: 'https://example.test' }); + const oldLogin = client.login('alice', 'password'); + await client.login('bob', 'password'); + first.resolve(Response.json({ access_token: 'alice', refresh_token: 'alice-refresh' })); + await expect(oldLogin).rejects.toThrow('superseded'); + expect(client.getTokens()).toEqual({ accessToken: 'new-access', refreshToken: 'new-refresh' }); +}); + +test.each(['login', 'register'] as const)('failed %s does not discard an in-flight refresh of the current account', async (method) => { + const refresh = deferred(); + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return refresh.promise; + if (url.endsWith('/auth/login') || url.endsWith('/auth/register')) return unauthorized(); + return bearer(init) === 'Bearer new-access' ? Response.json({ user_id: 'alice' }) : unauthorized(); + })); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const me = client.getMe(); + await vi.waitFor(() => expect(vi.mocked(fetch)).toHaveBeenCalledTimes(2)); + await expect(client[method]('bob', 'bad-password')).rejects.toMatchObject({ status: 401 }); + refresh.resolve(Response.json(tokens)); + await expect(me).resolves.toEqual({ user_id: 'alice' }); + expect(client.getTokens()).toEqual({ accessToken: 'new-access', refreshToken: 'new-refresh' }); +}); + +test.each([ + { authorization: 'Bearer external' }, + new Headers({ AUTHORIZATION: 'Bearer external' }), + { forEach: (callback: (value: string, key: string) => void) => callback('Bearer external', 'authorization') } as Headers, +])('explicit authorization is not refreshed using client credentials', async (headers) => { + const fetch = vi.fn(async () => unauthorized()); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const response = await client.fetchResponse('/resource', { headers }); + expect(response.status).toBe(401); + expect(fetch).toHaveBeenCalledTimes(1); + expect(client.getTokens()).toEqual({ accessToken: 'old', refreshToken: 'old-refresh' }); +}); + +test('raw mutations require explicit retry and a replayable body', async () => { + let refreshes = 0; + const fetch = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { refreshes++; return Response.json(tokens); } + return bearer(init) === 'Bearer new-access' ? new Response('ok') : unauthorized(); + }); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + expect((await client.fetchResponse('/write', { method: 'POST', body: '{}' })).status).toBe(401); + expect(refreshes).toBe(0); + expect((await client.fetchResponse('/write', { method: 'POST', body: '{}' }, { retryUnauthorized: true })).status).toBe(200); + expect(refreshes).toBe(1); + client.setTokens('old', 'old-refresh'); + const body = new ReadableStream({ start(controller) { controller.close(); } }); + expect((await client.fetchResponse('/write', { method: 'POST', body }, { retryUnauthorized: true })).status).toBe(401); + expect(refreshes).toBe(1); +}); + +test.each(['chat', 'work'] as const)('%s stream and REST share refresh without replaying an accepted stream', async (kind) => { + const refresh = deferred(); + let refreshes = 0; + const fetch = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { refreshes++; return refresh.promise; } + if (bearer(init) === 'Bearer old') return unauthorized(); + if (url.endsWith('/auth/me')) return Response.json({ user_id: 'alice' }); + return new Response('data: {"type":"run_finished","run_id":"run-1","status":"completed"}\n\n'); + }); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const events: unknown[] = []; + if (kind === 'chat') client.streamChat({ message: 'hi', modelSelection: { offeringId: 'offer' } }, { onEvent: event => events.push(event) }); + else client.continueWorkBranch('work-1', 'branch-1', { requestId: 'request-1', attachmentId: 'attachment-1', message: 'hi' }, { onEvent: event => events.push(event) }); + const rest = client.getMe(); + await vi.waitFor(() => expect(refreshes).toBe(1)); + refresh.resolve(Response.json(tokens)); + await expect(rest).resolves.toEqual({ user_id: 'alice' }); + await vi.waitFor(() => expect(events).toHaveLength(1)); + expect(events[0]).toMatchObject({ type: 'run_finished' }); + expect(refreshes).toBe(1); + expect(fetch.mock.calls.filter(([url]) => url.includes(kind === 'chat' ? '/chat/stream' : '/turns'))).toHaveLength(2); +}); + +test('an aborted refresh waiter settles while other callers keep sharing the refresh', async () => { + const refresh = deferred(); + const fetch = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return refresh.promise; + return bearer(init) === 'Bearer new-access' ? new Response('ok') : unauthorized(); + }); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const controller = new AbortController(); + const aborted = client.fetchResponse('/cancelled', { signal: controller.signal }); + const other = client.fetchResponse('/other'); + await vi.waitFor(() => expect(fetch.mock.calls.filter(([url]) => url.endsWith('/auth/refresh'))).toHaveLength(1)); + controller.abort(); + // This must settle without releasing the pending shared refresh. + await expect(aborted).rejects.toMatchObject({ name: 'AbortError' }); + expect(fetch.mock.calls.filter(([url]) => url.endsWith('/cancelled'))).toHaveLength(1); + refresh.resolve(Response.json(tokens)); + expect((await other).status).toBe(200); + expect(fetch.mock.calls.filter(([url]) => url.endsWith('/auth/refresh'))).toHaveLength(1); +}); + +test('an abort between refresh completion and replay preserves the caller reason', async () => { + const controller = new AbortController(); + const reason = new Error('caller stopped'); + const fetch = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return Response.json(tokens); + if (bearer(init) === 'Bearer new-access') { + if (url.endsWith('/first')) controller.abort(reason); + return new Response('ok'); + } + return unauthorized(); + }); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const first = client.fetchResponse('/first'); + const second = client.fetchResponse('/second', { signal: controller.signal }); + await expect(first).resolves.toMatchObject({ status: 200 }); + await expect(second).rejects.toBe(reason); + expect(fetch.mock.calls.filter(([url]) => url.endsWith('/second'))).toHaveLength(1); +}); + +test('aborted waiters release their registrations even when refresh never settles', async () => { + const refresh = deferred(); + vi.stubGlobal('fetch', vi.fn(async (url: string) => url.endsWith('/auth/refresh') ? refresh.promise : unauthorized())); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const reason = new Error('caller stopped'); + for (let i = 0; i < 20; i++) { + const controller = new AbortController(); + const removed = vi.spyOn(controller.signal, 'removeEventListener'); + const request = client.fetchResponse('/resource', { signal: controller.signal }); + await vi.waitFor(() => expect(vi.mocked(fetch)).toHaveBeenCalledTimes(i + 2)); + controller.abort(reason); + await expect(request).rejects.toBe(reason); + expect(removed).toHaveBeenCalledWith('abort', expect.any(Function)); + } + const remaining = client.fetchResponse('/remaining'); + await vi.waitFor(() => expect(vi.mocked(fetch)).toHaveBeenCalledTimes(22)); + refresh.resolve(Response.json(tokens)); + expect((await remaining).status).toBe(401); + expect(vi.mocked(fetch).mock.calls.filter(([url]) => String(url).endsWith('/auth/refresh'))).toHaveLength(1); +}); + +test('closing a stream while refresh is pending prevents a second POST', async () => { + const refresh = deferred(); + const fetch = vi.fn(async (url: string) => url.endsWith('/auth/refresh') ? refresh.promise : unauthorized()); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const stream = client.streamChat({ message: 'hi', modelSelection: { offeringId: 'offer' } }, { onEvent: () => {} }); + await vi.waitFor(() => expect(fetch).toHaveBeenCalledTimes(2)); + stream.close(); + refresh.resolve(Response.json(tokens)); + await vi.waitFor(() => expect(client.getTokens().accessToken).toBe('new-access')); + expect(fetch.mock.calls.filter(([url]) => url.endsWith('/chat/stream'))).toHaveLength(1); +}); + +test('reconnecting an old stream cannot issue a request under a replacement account', async () => { + const fetch = vi.fn(async () => new Response('data: {"type":"run_finished","status":"completed"}\n\n')); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'alice' }); + const events: unknown[] = []; + const stream = client.streamChat({ message: 'hi', modelSelection: { offeringId: 'offer' } }, { onEvent: event => events.push(event) }); + await vi.waitFor(() => expect(events).toHaveLength(1)); + client.setTokens('bob'); + await stream.connect(); + expect(fetch).toHaveBeenCalledTimes(1); + expect(events[1]).toMatchObject({ type: 'error', retryable: false }); +}); + +test('stream requests preserve configured headers while using the current owned bearer', async () => { + const fetch = vi.fn(async (_url: string, _init?: RequestInit) => new Response('data: {"type":"run_finished","status":"completed"}\n\n')); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'current', headers: { + Authorization: 'Bearer configured', Accept: 'application/custom', 'Content-Type': 'application/custom+json', + } }); + const events: unknown[] = []; + client.streamChat({ message: 'hi', modelSelection: { offeringId: 'offer' } }, { onEvent: event => events.push(event) }); + await vi.waitFor(() => expect(events).toHaveLength(1)); + const headers = new Headers(fetch.mock.calls[0][1]?.headers); + expect(headers.get('authorization')).toBe('Bearer current'); + expect(headers.get('accept')).toBe('application/custom'); + expect(headers.get('content-type')).toBe('application/custom+json'); +}); diff --git a/packages/sdk/src/__tests__/client-edge-stream.test.ts b/packages/sdk/src/__tests__/client-edge-stream.test.ts index cbf59ddaad..7c2cd2ff11 100644 --- a/packages/sdk/src/__tests__/client-edge-stream.test.ts +++ b/packages/sdk/src/__tests__/client-edge-stream.test.ts @@ -63,8 +63,8 @@ describe('AstraClient — streamChat', () => { const [url, init] = (globalThis.fetch as ReturnType).mock.calls[0] as [string, RequestInit]; expect(String(url).endsWith(PATH_CHAT_STREAM)).toBe(true); expect(init.method).toBe('POST'); - const headers = init.headers as Record; - expect(headers['Authorization']).toBe('Bearer tok-1'); + const headers = new Headers(init.headers); + expect(headers.get('Authorization')).toBe('Bearer tok-1'); expect(init.body).toBe( JSON.stringify( chatRequestToWire({ diff --git a/packages/sdk/src/__tests__/client.test.ts b/packages/sdk/src/__tests__/client.test.ts index 944c9c70dd..ebcaa4f1c7 100644 --- a/packages/sdk/src/__tests__/client.test.ts +++ b/packages/sdk/src/__tests__/client.test.ts @@ -428,11 +428,30 @@ describe("AstraClient — Runs", () => { 'data: {"type":"turn_complete"}\n\n'; globalThis.fetch = mockFetch(200, sse); - const events = await createClient().getRunEvents("r1", 5); + const client = new AstraClient({ + baseUrl: "http://localhost:8000", pathPrefix: "/api", accessToken: "current-token", + headers: { accept: "application/json", AUTHORIZATION: "Bearer stale-token", "X-Custom": "preserved" }, + }); + const events = await client.getRunEvents("r1", 5); expect(events).toHaveLength(2); expect( (globalThis.fetch as ReturnType).mock.calls[0][0], - ).toContain("last_index=5"); + ).toBe("http://localhost:8000/api/chat/runs/r1/stream?last_index=5"); + const headers = new Headers((globalThis.fetch as ReturnType).mock.calls[0][1].headers); + expect(headers.get("Accept")).toBe("text/event-stream"); + expect(headers.get("Authorization")).toBe("Bearer current-token"); + expect(headers.get("X-Custom")).toBe("preserved"); + }); + + test("getRunEvents preserves structured API errors", async () => { + globalThis.fetch = mockFetch(429, { + detail: "Try later", code: "rate_limited", category: "capacity", + retryable: true, action_hints: ["retry_later"], + }, { "Content-Type": "application/json" }); + await expect(createClient().getRunEvents("r1")).rejects.toMatchObject({ + status: 429, body: "Try later", path: "/chat/runs/r1/stream?last_index=0", + code: "rate_limited", category: "capacity", retryable: true, actionHints: ["retry_later"], + }); }); test("getSessionConversation preserves typed input identity and raw cursors", async () => { diff --git a/packages/sdk/src/__tests__/raw-response.test.ts b/packages/sdk/src/__tests__/raw-response.test.ts new file mode 100644 index 0000000000..231d74170f --- /dev/null +++ b/packages/sdk/src/__tests__/raw-response.test.ts @@ -0,0 +1,44 @@ +import { AstraClient } from '../client'; + +afterEach(() => vi.unstubAllGlobals()); + +test('raw transport leaves error bodies and FormData untouched; typed fetch keeps JSON defaults', async () => { + const fetchMock = vi.fn(async (_url: string, _init?: RequestInit) => new Response('upstream detail', { status: 403 })); + vi.stubGlobal('fetch', fetchMock); + const client = new AstraClient({ baseUrl: 'https://example.test/base', pathPrefix: '/api' }); + const body = new FormData(); + body.set('file', new Blob(['contents']), 'test.txt'); + const signal = new AbortController().signal; + const response = await client.fetchResponse('/upload', { method: 'POST', body, signal }); + expect(response.status).toBe(403); + expect(response.bodyUsed).toBe(false); + expect(await response.text()).toBe('upstream detail'); + expect(fetchMock.mock.calls[0]).toEqual(['https://example.test/base/api/upload', expect.objectContaining({ body, signal })]); + expect(new Headers(fetchMock.mock.calls[0][1]?.headers).has('Content-Type')).toBe(false); + await expect(client.fetch('/typed', { method: 'POST', body: '{}' })).rejects.toMatchObject({ status: 403 }); + expect(new Headers(fetchMock.mock.calls[1][1]?.headers).get('Content-Type')).toBe('application/json'); +}); + +test('anonymous raw requests suppress configured auth and refresh but allow explicit headers', async () => { + const fetchMock = vi.fn(async (_url: string, _init?: RequestInit) => new Response('unauthorized', { status: 401 })); + vi.stubGlobal('fetch', fetchMock); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'access', refreshToken: 'refresh', headers: { authorization: 'Bearer configured' } }); + await client.fetchResponse('/public', {}, { auth: 'none' }); + await client.fetchResponse('/public', { headers: { Authorization: 'Bearer explicit' } }, { auth: 'none' }); + expect(fetchMock).toHaveBeenCalledTimes(2); + expect(new Headers(fetchMock.mock.calls[0][1]?.headers).has('Authorization')).toBe(false); + expect(new Headers(fetchMock.mock.calls[1][1]?.headers).get('Authorization')).toBe('Bearer explicit'); +}); + +test.each([null, {}, { access_token: 'new', refresh_token: 1 }])('malformed refresh does not publish credentials or notify: %j', async payload => { + const notify = vi.fn(); + const fetchMock = vi.fn(async (url: string) => url.endsWith('/auth/refresh') + ? Response.json(payload) : new Response('expired', { status: 401 })); + vi.stubGlobal('fetch', fetchMock); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'refresh', onTokenRefresh: notify }); + const snapshot = client.getTokens(); + await expect(client.getMe()).rejects.toMatchObject({ status: 401 }); + expect(client.getTokens()).toEqual(snapshot); + expect(notify).not.toHaveBeenCalled(); + expect(fetchMock).toHaveBeenCalledTimes(2); +}); diff --git a/packages/sdk/src/__tests__/token-refresh-concurrency.test.ts b/packages/sdk/src/__tests__/token-refresh-concurrency.test.ts new file mode 100644 index 0000000000..a5f4bee567 --- /dev/null +++ b/packages/sdk/src/__tests__/token-refresh-concurrency.test.ts @@ -0,0 +1,481 @@ +import { AstraClient } from '../client'; + +function deferred() { + let resolve!: (value: T) => void; + const promise = new Promise((done) => { resolve = done; }); + return { promise, resolve }; +} + +const rotated = { access_token: 'new-access', refresh_token: 'new-refresh' }; +const unauthorized = () => Response.json({ detail: 'expired' }, { status: 401 }); +const authorization = (init?: RequestInit) => new Headers(init?.headers).get('Authorization'); + +afterEach(() => vi.unstubAllGlobals()); + +test('an asynchronous reentrant notification cannot rotate forever on persistent 401', async () => { + let refreshes = 0; + let recovered = false; + vi.stubGlobal('fetch', vi.fn(async (url: string) => { + if (url.endsWith('/auth/refresh')) { + refreshes++; + // Keep a broken implementation bounded during the regression test. + return refreshes > 4 ? new Response('stop', { status: 503 }) + : Response.json({ access_token: `access-${refreshes}`, refresh_token: `refresh-${refreshes}` }); + } + return recovered ? Response.json({ user_id: 'alice' }) : unauthorized(); + })); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + onTokenRefresh: async () => { await Promise.resolve(); await client.getMe(); }, + }); + await expect(client.getMe()).rejects.toMatchObject({ status: 401 }); + expect(refreshes).toBe(1); + recovered = true; + await expect(client.getMe()).resolves.toEqual({ user_id: 'alice' }); +}); + +test.each(['pending', 'failed'] as const)('a delayed old-token 401 cannot bypass a %s notification', async (when) => { + const late401 = deferred(); + const notification = deferred(); + const entered = deferred(); + const calls: string[] = []; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return Response.json(rotated); + calls.push(`${url.split('/').pop()}:${authorization(init)}`); + if (url.endsWith('/late') && authorization(init) === 'Bearer old-access') return late401.promise; + return authorization(init) === 'Bearer old-access' ? unauthorized() : new Response('ok'); + })); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + onTokenRefresh: async () => { entered.resolve(); await notification.promise; throw new Error('persistence failed'); }, + }); + const late = client.fetchResponse('/late'); + const first = client.fetchResponse('/first'); + await entered.promise; + if (when === 'failed') notification.resolve(); + if (when === 'failed') expect((await first).status).toBe(401); + late401.resolve(unauthorized()); + expect((await late).status).toBe(401); + expect(calls.filter(call => call.startsWith('late:'))).toEqual(['late:Bearer old-access']); + if (when === 'pending') notification.resolve(); + expect((await first).status).toBe(401); +}); + +test('a callback awaiting a delayed old-token response does not deadlock', async () => { + const late401 = deferred(); + const entered = deferred(); + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return Response.json(rotated); + if (url.endsWith('/late') && authorization(init) === 'Bearer old-access') return late401.promise; + return authorization(init) === 'Bearer old-access' ? unauthorized() : new Response('ok'); + })); + let late!: Promise; + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + onTokenRefresh: async () => { entered.resolve(); expect((await late).status).toBe(401); }, + }); + late = client.fetchResponse('/late'); + const first = client.fetchResponse('/first'); + await entered.promise; + late401.resolve(unauthorized()); + expect((await first).status).toBe(200); +}); + +test('REST and buffered SSE share one refresh of a single-use token', async () => { + const refresh = deferred(); + const onTokenRefresh = vi.fn(); + let refreshCount = 0; + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + refreshCount++; + expect(JSON.parse(String(init?.body))).toEqual({ refresh_token: 'old-refresh' }); + // The runtime revokes the token on its first successful refresh. + return refreshCount === 1 ? refresh.promise : unauthorized(); + } + if (authorization(init) === 'Bearer old-access') return unauthorized(); + expect(authorization(init)).toBe('Bearer new-access'); + return url.includes('/stream') + ? new Response('data: {"type":"run_finished"}\n\n') + : Response.json({ user_id: 'alice' }); + }); + vi.stubGlobal('fetch', fetchMock); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', + refreshToken: 'old-refresh', onTokenRefresh, + }); + const results = Promise.allSettled([client.getMe(), client.getRunEvents('run-1')]); + await vi.waitFor(() => expect(refreshCount).toBeGreaterThan(0)); + refresh.resolve(Response.json(rotated)); + expect((await results).map((result) => result.status)).toEqual(['fulfilled', 'fulfilled']); + expect(refreshCount).toBe(1); + expect(onTokenRefresh).toHaveBeenCalledTimes(1); +}); + +test('refresh-token-only clients resume a session with one shared rotation', async () => { + const refresh = deferred(); + let rotations = 0; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { rotations++; return refresh.promise; } + return authorization(init) === 'Bearer new-access' + ? Response.json({ user_id: 'alice' }) : unauthorized(); + })); + const client = new AstraClient({ baseUrl: 'https://example.test', refreshToken: 'old-refresh' }); + const results = Promise.all([client.getMe(), client.getMe()]); + await vi.waitFor(() => expect(rotations).toBe(1)); + refresh.resolve(Response.json(rotated)); + expect(await results).toEqual([{ user_id: 'alice' }, { user_id: 'alice' }]); + expect(rotations).toBe(1); +}); + +test.each(['anonymous', 'request authorization', 'configured authorization'] as const)( + 'a refresh-token-only client does not refresh %s requests', async (mode) => { + const fetch = vi.fn(async () => unauthorized()); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ + baseUrl: 'https://example.test', refreshToken: 'old-refresh', + headers: mode === 'configured authorization' ? { Authorization: 'Bearer external' } : undefined, + }); + const response = await client.fetchResponse('/resource', + mode === 'request authorization' ? { headers: { Authorization: 'Bearer external' } } : undefined, + mode === 'anonymous' ? { auth: 'none' } : undefined); + expect(response.status).toBe(401); + expect(fetch).toHaveBeenCalledTimes(1); + }, +); + +test('a late 401 reuses the token already rotated by another request', async () => { + const late = deferred(); + let initialRequests = 0; + let refreshCount = 0; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + refreshCount++; + return Response.json(rotated); + } + if (authorization(init) === 'Bearer old-access') { + return ++initialRequests === 1 ? late.promise : unauthorized(); + } + return Response.json({ user_id: 'alice' }); + })); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + }); + const first = client.getMe(); + await client.getMe(); + late.resolve(unauthorized()); + await expect(first).resolves.toEqual({ user_id: 'alice' }); + expect(refreshCount).toBe(1); +}); + +test('a late 401 joins a newer refresh instead of retrying its rejected token', async () => { + const late = deferred(); + const secondRefresh = deferred(); + const secondStarted = deferred(); + const lateTokens: string[] = []; + let refreshCount = 0; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + if (++refreshCount === 1) return Response.json(rotated); + secondStarted.resolve(); + return secondRefresh.promise; + } + const token = authorization(init); + if (url.endsWith('/late')) { + lateTokens.push(token ?? ''); + if (token === 'Bearer old-access') return late.promise; + } + return token === 'Bearer final-access' || (url.endsWith('/first') && token === 'Bearer new-access') + ? new Response('ok') : unauthorized(); + })); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + }); + const delayed = client.fetchResponse('/late'); + expect((await client.fetchResponse('/first')).status).toBe(200); + const next = client.fetchResponse('/next'); + await secondStarted.promise; + late.resolve(unauthorized()); + await new Promise(resolve => setTimeout(resolve, 0)); + secondRefresh.resolve(Response.json({ access_token: 'final-access', refresh_token: 'final-refresh' })); + expect((await delayed).status).toBe(200); + expect((await next).status).toBe(200); + expect(lateTokens).toEqual(['Bearer old-access', 'Bearer final-access']); + expect(refreshCount).toBe(2); +}); + +test('a delayed 401 can use a later successfully persisted recovery', async () => { + const late = deferred(); + let rotations = 0; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + rotations++; + return Response.json(rotations === 1 ? rotated + : { access_token: 'recovered-access', refresh_token: 'recovered-refresh' }); + } + if (url.endsWith('/late') && authorization(init) === 'Bearer old-access') return late.promise; + return authorization(init) === 'Bearer recovered-access' ? new Response('ok') : unauthorized(); + })); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + onTokenRefresh: async () => { if (rotations === 1) throw new Error('persistence failed'); }, + }); + const delayed = client.fetchResponse('/late'); + expect((await client.fetchResponse('/first')).status).toBe(401); + expect((await client.fetchResponse('/recovery')).status).toBe(200); + late.resolve(unauthorized()); + expect((await delayed).status).toBe(200); + expect(rotations).toBe(2); +}); + +test.each(['success', 'failure'] as const)( + 'a delayed 401 joins a %s recovery after an earlier notification failure', async (result) => { + const late = deferred(); + const recovery = deferred(); + const recoveryStarted = deferred(); + let rotations = 0; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + if (++rotations === 1) return Response.json(rotated); + recoveryStarted.resolve(); + return recovery.promise; + } + if (url.endsWith('/late') && authorization(init) === 'Bearer old-access') return late.promise; + return authorization(init) === 'Bearer recovered-access' ? new Response('ok') : unauthorized(); + })); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + onTokenRefresh: async () => { if (rotations === 1) throw new Error('persistence failed'); }, + }); + const delayed = client.fetchResponse('/late'); + expect((await client.fetchResponse('/first')).status).toBe(401); + const recovering = client.fetchResponse('/recovery'); + await recoveryStarted.promise; + late.resolve(unauthorized()); + await new Promise(resolve => setTimeout(resolve, 0)); + recovery.resolve(result === 'success' + ? Response.json({ access_token: 'recovered-access', refresh_token: 'recovered-refresh' }) + : new Response('unavailable', { status: 503 })); + expect((await recovering).status).toBe(result === 'success' ? 200 : 401); + expect((await delayed).status).toBe(result === 'success' ? 200 : 401); + expect(rotations).toBe(2); + }, +); + +test.each(['http', 'network'] as const)('failed %s refresh is shared but a later attempt can recover', async (failure) => { + const refresh = deferred(); + let refreshCount = 0; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + refreshCount++; + if (refreshCount > 1) return Response.json(rotated); + await refresh.promise; + if (failure === 'network') throw new TypeError('connection lost'); + return unauthorized(); + } + return authorization(init) === 'Bearer old-access' + ? unauthorized() : Response.json({ user_id: 'alice' }); + })); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + }); + const results = Promise.allSettled([client.getMe(), client.getMe()]); + await vi.waitFor(() => expect(refreshCount).toBeGreaterThan(0)); + refresh.resolve(unauthorized()); + expect((await results).map((result) => result.status)).toEqual(['rejected', 'rejected']); + expect(refreshCount).toBe(1); + await expect(client.getMe()).resolves.toEqual({ user_id: 'alice' }); + expect(refreshCount).toBe(2); +}); + +test.each(['setTokens', 'login', 'register', 'logout'] as const)( + 'a refresh started before %s cannot restore the previous credentials', async (action) => { + const refresh = deferred(); + const onTokenRefresh = vi.fn(); + let refreshCount = 0; + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { refreshCount++; return refresh.promise; } + if (url.endsWith('/auth/logout')) return new Response(null, { status: 204 }); + if (url.endsWith('/auth/login') || url.endsWith('/auth/register')) { + return Response.json({ access_token: 'bob-access', refresh_token: 'bob-refresh' }); + } + return authorization(init) === 'Bearer old-access' + ? unauthorized() : Response.json({ user_id: 'bob' }); + }); + vi.stubGlobal('fetch', fetchMock); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', + refreshToken: 'old-refresh', onTokenRefresh, + }); + const pending = Promise.allSettled([client.getMe()]); + await vi.waitFor(() => expect(refreshCount).toBe(1)); + if (action === 'setTokens') client.setTokens('bob-access', 'bob-refresh'); + else if (action === 'logout') await client.logout(); + else await client[action]('bob', 'password'); + const callsBeforeCompletion = fetchMock.mock.calls.length; + refresh.resolve(Response.json(rotated)); + expect((await pending)[0].status).toBe('rejected'); + expect(fetchMock).toHaveBeenCalledTimes(callsBeforeCompletion); + expect(onTokenRefresh).not.toHaveBeenCalled(); + await client.getMe(); + expect(authorization(fetchMock.mock.lastCall?.[1])).toBe( + action === 'logout' ? null : 'Bearer bob-access', + ); + }, +); + +test('an async notification can receive 401 without blocking its callers or later recovery', async () => { + const refreshTokens: string[] = []; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + refreshTokens.push(JSON.parse(String(init?.body)).refresh_token); + return Response.json(refreshTokens.length === 1 ? rotated + : { access_token: 'recovered-access', refresh_token: 'recovered-refresh' }); + } + return authorization(init) === 'Bearer recovered-access' + ? Response.json({ user_id: 'alice' }) : unauthorized(); + })); + const onTokenRefresh = vi.fn(async () => { await client.getMe(); }); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', + refreshToken: 'old-refresh', onTokenRefresh, + }); + const results = await Promise.allSettled([client.getMe(), client.getRunEvents('run-1')]); + expect(results.map(result => result.status)).toEqual(['rejected', 'rejected']); + expect(refreshTokens).toEqual(['old-refresh']); + expect(onTokenRefresh).toHaveBeenCalledTimes(1); + await expect(client.getMe()).resolves.toEqual({ user_id: 'alice' }); + expect(refreshTokens).toEqual(['old-refresh', 'new-refresh']); + expect(onTokenRefresh).toHaveBeenCalledTimes(2); +}); + +test('a callback 401 does not start an overlapping rotation', async () => { + const refreshTokens: string[] = []; + const callbackResults: unknown[] = []; + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + refreshTokens.push(JSON.parse(String(init?.body)).refresh_token); + return Response.json(refreshTokens.length === 1 ? rotated + : { access_token: 'recovered-access', refresh_token: 'recovered-refresh' }); + } + return url.endsWith('/auth/me') && authorization(init) === 'Bearer recovered-access' + ? Response.json({ user_id: 'alice' }) : unauthorized(); + }); + vi.stubGlobal('fetch', fetchMock); + const onTokenRefresh = vi.fn(async () => { callbackResults.push(await client.getMe().catch(error => error.status)); }); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', + refreshToken: 'old-refresh', onTokenRefresh, + }); + await expect(client.getRunEvents('run-1')).rejects.toMatchObject({ status: 401 }); + expect(refreshTokens).toEqual(['old-refresh']); + expect(callbackResults).toEqual([401]); + expect(onTokenRefresh).toHaveBeenCalledTimes(1); + const streamRequests = fetchMock.mock.calls.filter(([url]) => url.includes('/stream')); + expect(streamRequests.map(([, init]) => authorization(init))).toEqual([ + 'Bearer old-access', 'Bearer new-access', + ]); + expect(streamRequests.every(([, init]) => new Headers(init?.headers).get('Accept') === 'text/event-stream')).toBe(true); +}); + +test('already joined callers await an async notification that successfully reenters the client', async () => { + const notification = deferred(); + const entered = deferred(); + const settled = vi.fn(); + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return Response.json(rotated); + if (authorization(init) === 'Bearer old-access') return unauthorized(); + return url.includes('/stream') ? new Response('data: {"type":"run_finished"}\n\n') + : Response.json({ user_id: 'alice' }); + })); + const onTokenRefresh = vi.fn(async () => { + await client.getMe(); + entered.resolve(); + await notification.promise; + }); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', + refreshToken: 'old-refresh', onTokenRefresh, + }); + const results = Promise.allSettled([client.getMe(), client.getRunEvents('run-1')] + .map(request => request.then(settled, settled))); + await entered.promise; + expect(settled).not.toHaveBeenCalled(); + notification.resolve(); + await results; + expect(settled).toHaveBeenCalledTimes(2); + expect(settled).toHaveBeenCalledWith({ user_id: 'alice' }); + expect(settled).toHaveBeenCalledWith([{ type: 'run_finished' }]); + expect(onTokenRefresh).toHaveBeenCalledTimes(1); +}); + +test.each(['throw', 'reject'] as const)('a notification %s settles joined callers and preserves rotated credentials', async failure => { + const refresh = deferred(); + const refreshTokens: string[] = []; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + refreshTokens.push(JSON.parse(String(init?.body)).refresh_token); + return refreshTokens.length === 1 ? refresh.promise + : Response.json({ access_token: 'recovered-access', refresh_token: 'recovered-refresh' }); + } + return authorization(init) === 'Bearer recovered-access' + ? Response.json({ user_id: 'alice' }) : unauthorized(); + })); + const onTokenRefresh = vi.fn(() => { + if (onTokenRefresh.mock.calls.length > 1) return; + const error = new Error('notification failed'); + if (failure === 'throw') throw error; + return Promise.reject(error); + }); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', + refreshToken: 'old-refresh', onTokenRefresh, + }); + const results = Promise.allSettled([client.getMe(), client.getRunEvents('run-1')]); + await vi.waitFor(() => expect(refreshTokens).toHaveLength(1)); + refresh.resolve(Response.json(rotated)); + expect((await results).map(result => result.status)).toEqual(['rejected', 'rejected']); + await expect(client.getMe()).resolves.toEqual({ user_id: 'alice' }); + expect(refreshTokens).toEqual(['old-refresh', 'new-refresh']); + expect(onTokenRefresh).toHaveBeenCalledTimes(2); +}); + +test('an older notification cannot replay across accounts or release a newer refresh', async () => { + const notification = deferred(); + const notified = deferred(); + const bobRefresh = deferred(); + const bobRefreshing = deferred(); + const refreshTokens: string[] = []; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + const token = JSON.parse(String(init?.body)).refresh_token; + refreshTokens.push(token); + if (token === 'old-refresh') return Response.json(rotated); + bobRefreshing.resolve(); + return bobRefresh.promise; + } + return authorization(init) === 'Bearer bob-access' + ? Response.json({ user_id: 'bob' }) : unauthorized(); + })); + const onTokenRefresh = vi.fn(async ({ accessToken }: { accessToken: string }) => { + if (accessToken !== 'new-access') return; + notified.resolve(); + await notification.promise; + }); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', + refreshToken: 'old-refresh', onTokenRefresh, + }); + const alice = Promise.allSettled([client.getMe()]); + await notified.promise; + client.setTokens('bob-expired', 'bob-refresh'); + const firstBob = client.getMe(); + await bobRefreshing.promise; + notification.resolve(); + expect((await alice)[0].status).toBe('rejected'); + const secondBob = client.getMe(); + await new Promise(resolve => setImmediate(resolve)); + expect(refreshTokens).toEqual(['old-refresh', 'bob-refresh']); + bobRefresh.resolve(Response.json({ access_token: 'bob-access', refresh_token: 'bob-new-refresh' })); + expect(await Promise.all([firstBob, secondBob])).toEqual([{ user_id: 'bob' }, { user_id: 'bob' }]); + expect(onTokenRefresh).toHaveBeenCalledTimes(2); +}); diff --git a/packages/sdk/src/client.ts b/packages/sdk/src/client.ts index 02b78c224f..5efeb6f6e6 100644 --- a/packages/sdk/src/client.ts +++ b/packages/sdk/src/client.ts @@ -242,7 +242,6 @@ import { SSEClient, parseSseDataEvents } from "./sse-client"; import { headersInitToRecord, readAstraError, - readAstraErrorDetail, } from "./http"; import { modelSelectionToWire } from "./wire"; import { @@ -623,15 +622,32 @@ export function chatRequestToWire(req: ChatRequest): Record { * same layout as `astra-thin-client` / `astra-server`; set `pathPrefix` when a * gateway mounts the API under a prefix. */ +type RefreshAttempt = { + phase: "refreshing" | "notifying" | "done"; + result: boolean; + waiters: Set<(result: boolean) => void>; +}; + +type TokenRevision = { + accessToken: string | null; + refreshToken: string | null; + notification: "pending" | "succeeded" | "failed"; + attempt: RefreshAttempt | null; +}; + +function tokenRevision(accessToken: string | null, refreshToken: string | null): TokenRevision { + return { accessToken, refreshToken, notification: "succeeded", attempt: null }; +} + export class AstraClient { private config: AstraClientConfig; - private accessToken: string | null; - private refreshTokenValue: string | null; + private tokens: TokenRevision; + private authGeneration = 0; + private authAttempt: object | null = null; constructor(config: AstraClientConfig) { this.config = config; - this.accessToken = config.accessToken ?? null; - this.refreshTokenValue = config.refreshToken ?? null; + this.tokens = tokenRevision(config.accessToken ?? null, config.refreshToken ?? null); } private apiPath(path: string): string { @@ -653,34 +669,47 @@ export class AstraClient { const email = options?.email ?? `${username}@users.local.astra`; const body: Record = { username, email, password }; if (options?.displayName) body.display_name = options.displayName; - const result = await this.post(PATH_AUTH_REGISTER, body); - this.accessToken = result.access_token; - this.refreshTokenValue = result.refresh_token; - return result; + return this.authenticate(PATH_AUTH_REGISTER, body); } /** Log in with username/password. Stores tokens automatically. */ async login(username: string, password: string): Promise { - const result = await this.post(PATH_AUTH_LOGIN, { + return this.authenticate(PATH_AUTH_LOGIN, { username, password, }); - this.accessToken = result.access_token; - this.refreshTokenValue = result.refresh_token; - return result; } - /** Log out and clear stored tokens. Requires refresh token in client state. */ - async logout(): Promise { + private async authenticate(path: string, body: unknown): Promise { + const attempt = {}; + this.authAttempt = attempt; try { - if (this.refreshTokenValue) { - await this.post(PATH_AUTH_LOGOUT, { - refresh_token: this.refreshTokenValue, - }); + const response = await this.request(path, { + method: "POST", + body: JSON.stringify(body), + }, { auth: "none" }); + const result = await response.json() as AuthResult; + if (this.authAttempt !== attempt) { + throw new Error("Authentication attempt was superseded"); } + this.setTokens(result.access_token, result.refresh_token); + return result; } finally { - this.accessToken = null; - this.refreshTokenValue = null; + if (this.authAttempt === attempt) this.authAttempt = null; + } + } + + /** Log out and clear stored tokens. Requires refresh token in client state. */ + async logout(): Promise { + this.authGeneration++; + this.authAttempt = null; + const refreshToken = this.tokens.refreshToken; + this.tokens = tokenRevision(null, null); + if (refreshToken) { + await this.request(PATH_AUTH_LOGOUT, { + method: "POST", + body: JSON.stringify({ refresh_token: refreshToken }), + }, { auth: "none" }); } } @@ -743,72 +772,158 @@ export class AstraClient { } setTokens(accessToken: string, refreshToken?: string): void { - this.accessToken = accessToken; - if (refreshToken) this.refreshTokenValue = refreshToken; + this.authGeneration++; + this.authAttempt = null; + this.tokens = tokenRevision(accessToken, refreshToken || null); + } + + /** A detached snapshot of the credentials currently owned by this client. */ + getTokens(): Readonly<{ accessToken: string | null; refreshToken: string | null }> { + return { accessToken: this.tokens.accessToken, refreshToken: this.tokens.refreshToken }; + } + + private tryRefreshToken(generation: number, revision: TokenRevision, signal?: AbortSignal | null): Promise { + if (signal?.aborted) return Promise.reject(signal.reason); + if (generation !== this.authGeneration) return Promise.resolve(false); + // Do not overlap notifications: an older async persistence callback could + // otherwise overwrite the successor refresh token after rotation. + if (this.tokens.notification === "pending") return Promise.resolve(false); + let attempt = revision.attempt; + if (attempt?.phase === "notifying") return Promise.resolve(false); + if (attempt?.phase === "done") { + const newerAttempt = this.tokens.attempt; + if (newerAttempt?.phase === "refreshing") { + // Join the newer rotation instead of retrying a token that another + // request has already rejected. + attempt = newerAttempt; + } else { + return Promise.resolve( + (attempt.result || revision !== this.tokens) && this.tokens.notification === "succeeded", + ); + } + } + if (!attempt) { + if (revision !== this.tokens || !revision.refreshToken) return Promise.resolve(false); + attempt = { phase: "refreshing", result: false, waiters: new Set() }; + revision.attempt = attempt; + const current = attempt; + const finish = (result: boolean) => { + current.phase = "done"; + current.result = result; + // Failed network/HTTP refreshes may be tried again later. + if (revision === this.tokens) revision.attempt = null; + for (const settle of current.waiters) settle(result); + }; + void this.refreshTokens(generation, revision, current).then( + result => { + finish(result); + }, + () => { + finish(false); + }, + ); + } + const current = attempt; + return new Promise((resolve, reject) => { + const cleanup = () => { + current.waiters.delete(settle); + signal?.removeEventListener("abort", abort); + }; + const settle = (result: boolean) => { cleanup(); resolve(result); }; + const abort = () => { cleanup(); reject(signal?.reason); }; + current.waiters.add(settle); + signal?.addEventListener("abort", abort, { once: true }); + if (signal?.aborted) abort(); + }); } - private async tryRefreshToken(): Promise { - if (!this.refreshTokenValue) return false; + private async refreshTokens(generation: number, revision: TokenRevision, attempt: RefreshAttempt): Promise { try { const res = await fetch(this.apiPath(PATH_AUTH_REFRESH), { method: "POST", headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ refresh_token: this.refreshTokenValue }), + body: JSON.stringify({ refresh_token: revision.refreshToken }), + cache: "no-store", }); if (!res.ok) return false; const data = (await res.json()) as AuthResult; - this.accessToken = data.access_token; - this.refreshTokenValue = data.refresh_token; - await this.config.onTokenRefresh?.({ - accessToken: data.access_token, - refreshToken: data.refresh_token, - }); - return true; + if ( + generation !== this.authGeneration || + !data || + typeof data.access_token !== "string" || + typeof data.refresh_token !== "string" + ) return false; + const next = tokenRevision(data.access_token, data.refresh_token); + this.tokens = next; + if (this.config.onTokenRefresh) { + next.notification = "pending"; + attempt.phase = "notifying"; + try { + await this.config.onTokenRefresh({ accessToken: data.access_token, refreshToken: data.refresh_token }); + next.notification = "succeeded"; + } catch { + next.notification = "failed"; + return false; + } + } + return generation === this.authGeneration && this.tokens.notification === "succeeded"; } catch { return false; } } - private buildHeaders( - init?: RequestInit, - extra?: Record, - ): Record { - const headers = headersInitToRecord(this.config.headers ?? {}, extra); - const generated: Record = {}; - const method = (init?.method ?? "GET").toUpperCase(); - if ( - init?.body != null && - (method === "POST" || - method === "PUT" || - method === "PATCH" || - method === "DELETE") - ) { - generated["Content-Type"] = "application/json"; - } - if (this.accessToken) { - generated["Authorization"] = `Bearer ${this.accessToken}`; - } - return headersInitToRecord(headers, generated); + private buildHeaders(auth: "optional" | "none"): Record { + const headers = new Headers(headersInitToRecord(this.config.headers ?? {})); + if (auth === "none") headers.delete("Authorization"); + else if (this.tokens.accessToken) headers.set("Authorization", `Bearer ${this.tokens.accessToken}`); + return headersInitToRecord({}, headers); + } + + private streamHeaders(extra?: Record): Record { + const headers = new Headers(headersInitToRecord(this.config.headers ?? {})); + headers.delete("Authorization"); + return headersInitToRecord(headersInitToRecord({}, headers), extra); } // ─── HTTP helpers ────────────────────────────────────────────────── - private async request(path: string, init?: RequestInit): Promise { + /** Raw, body-neutral transport; returns the final response without consuming it. */ + async fetchResponse( + path: string, + init?: RequestInit, + options?: { auth?: "optional" | "none"; retryUnauthorized?: boolean }, + ): Promise { + const auth = options?.auth ?? "optional"; const url = this.apiPath(path); - let res = await fetch(url, { + const generation = this.authGeneration; + const revision = this.tokens; + const accessToken = revision.accessToken; + const requestHeaders = headersInitToRecord({}, init?.headers); + const explicitAuthorization = new Headers(requestHeaders).has("Authorization") + || (!accessToken && new Headers(headersInitToRecord(this.config.headers ?? {})).has("Authorization")); + const method = (init?.method ?? "GET").toUpperCase(); + const replayable = !(typeof ReadableStream !== "undefined" && init?.body instanceof ReadableStream); + const mayRetry = options?.retryUnauthorized ?? (method === "GET" || method === "HEAD"); + const send = () => fetch(url, { ...init, - headers: headersInitToRecord(this.buildHeaders(init), init?.headers), + headers: headersInitToRecord(this.buildHeaders(auth), requestHeaders), }); - - if (res.status === 401) { - const refreshed = await this.tryRefreshToken(); - if (refreshed) { - res = await fetch(url, { - ...init, - headers: headersInitToRecord(this.buildHeaders(init), init?.headers), - }); - } + let res = await send(); + if (res.status === 401 && auth !== "none" && (accessToken || revision.refreshToken) && !explicitAuthorization && mayRetry && replayable && !init?.signal?.aborted) { + const refreshed = await this.tryRefreshToken(generation, revision, init?.signal); + if (init?.signal?.aborted) throw init.signal.reason; + if (refreshed && generation === this.authGeneration && this.tokens.notification === "succeeded") res = await send(); } + return res; + } + + private async request(path: string, init?: RequestInit, options?: { auth?: "optional" | "none" }): Promise { + const method = (init?.method ?? "GET").toUpperCase(); + const json = init?.body != null && ["POST", "PUT", "PATCH", "DELETE"].includes(method); + const res = await this.fetchResponse(path, { + ...init, + headers: headersInitToRecord(json ? { "Content-Type": "application/json" } : {}, init?.headers), + }, { ...options, retryUnauthorized: true }); if (!res.ok) { const error = await readAstraError(res); @@ -2665,6 +2780,7 @@ export class AstraClient { ); } const url = this.apiPath(workBranchTurnsPath(workId, branchId)); + const generation = this.authGeneration; let client: SSEClient; client = new SSEClient({ url, @@ -2674,11 +2790,13 @@ export class AstraClient { attachment_id: input.attachmentId, message: input.message, }), - token: this.accessToken ?? undefined, - headers: { - ...this.config.headers, - [ASTRA_WORK_API_MAJOR_HEADER]: ASTRA_WORK_API_MAJOR, + transport: (_url, init) => { + if (generation !== this.authGeneration) throw new Error("Authentication changed before stream connection"); + return this.fetchResponse(workBranchTurnsPath(workId, branchId), init, { retryUnauthorized: true }); }, + headers: this.streamHeaders({ + [ASTRA_WORK_API_MAJOR_HEADER]: ASTRA_WORK_API_MAJOR, + }), onEvent: (event) => { try { callbacks.onEvent(decodeWorkTurnStreamEventV1(event)); @@ -3060,27 +3178,9 @@ export class AstraClient { */ async getRunEvents(runId: string, startIndex = 0): Promise { const path = `${chatRunStreamPath(runId)}?last_index=${encodeURIComponent(String(startIndex))}`; - const url = this.apiPath(path); - const streamHeaders = this.buildHeaders(undefined, { - Accept: "text/event-stream", + const res = await this.request(path, { + headers: { Accept: "text/event-stream" }, }); - let res = await fetch(url, { headers: streamHeaders }); - - if (res.status === 401) { - const refreshed = await this.tryRefreshToken(); - if (refreshed) { - res = await fetch(url, { - headers: this.buildHeaders(undefined, { - Accept: "text/event-stream", - }), - }); - } - } - - if (!res.ok) { - const body = await readAstraErrorDetail(res); - throw new AstraApiError(res.status, body, path); - } // Response.text() strips a BOM before the shared parser can apply its rule. const text = new TextDecoder("utf-8", { ignoreBOM: true }).decode(await res.arrayBuffer()); @@ -3398,12 +3498,16 @@ export class AstraClient { }, ): SSEClient { const url = this.apiPath(PATH_CHAT_STREAM); + const generation = this.authGeneration; const client = new SSEClient({ url, method: "POST", body: JSON.stringify(chatRequestToWire(request)), - token: this.accessToken ?? undefined, - headers: this.config.headers, + transport: (_url, init) => { + if (generation !== this.authGeneration) throw new Error("Authentication changed before stream connection"); + return this.fetchResponse(PATH_CHAT_STREAM, init, { retryUnauthorized: true }); + }, + headers: this.streamHeaders(), onEvent: callbacks.onEvent, onStateChange: callbacks.onStateChange, onRawLine: callbacks.onRawLine, diff --git a/packages/sdk/src/sse-client.ts b/packages/sdk/src/sse-client.ts index ba2d68692f..b4ee134931 100644 --- a/packages/sdk/src/sse-client.ts +++ b/packages/sdk/src/sse-client.ts @@ -187,7 +187,7 @@ export class SSEClient { headers = headersInitToRecord(headers, { 'Content-Type': 'application/json' }); } - const response = await fetch(this.options.url, { + const response = await (this.options.transport ?? fetch)(this.options.url, { method: this.options.method ?? 'GET', headers, body: this.options.body, diff --git a/packages/sdk/src/types.ts b/packages/sdk/src/types.ts index 593b31626e..f93bca4eaa 100644 --- a/packages/sdk/src/types.ts +++ b/packages/sdk/src/types.ts @@ -1209,6 +1209,8 @@ export type AstraClientConfig = { export type SSEClientOptions = { url: string; + /** Optional request transport for clients that own authentication. */ + transport?: (url: string, init: RequestInit) => Promise; token?: string; headers?: Record; onEvent: (event: StreamEvent) => void; diff --git a/scripts/schema/schema_inventory.py b/scripts/schema/schema_inventory.py index de16f64dde..137fdc094f 100755 --- a/scripts/schema/schema_inventory.py +++ b/scripts/schema/schema_inventory.py @@ -886,7 +886,7 @@ class AutoIncrementMetadata: "auth_refresh_tokens": TableMetadata( semantic_owner="astra_services::auth", state_class="durable refresh-session credential fact", - primary_query="refresh/logout lookup by token_hash, session_id, user_id, is_revoked, and expires_at", + primary_query="refresh compare-and-replace by token_hash and session; logout revokes by signed user/session claims", retention_policy="cleanup_expired_data prunes expired or revoked rows after refresh_token_days default 7 in ordered bounded batches; logout and refresh revoke specific token hashes", rebuildability="not rebuildable after token_hash/session binding is lost because refresh/logout idempotency and external session linkage depend on it", merge_guidance="keep separate from auth_users; refresh credentials are high-churn secrets with TTL/revocation lifecycle", diff --git a/web/__tests__/lib/chat-cold-recovery.test.ts b/web/__tests__/lib/chat-cold-recovery.test.ts index c9bdc72889..37dc2ce3eb 100644 --- a/web/__tests__/lib/chat-cold-recovery.test.ts +++ b/web/__tests__/lib/chat-cold-recovery.test.ts @@ -28,6 +28,7 @@ function setup() { globalThis.__astraWebStores = undefined; vi.mocked(requireRuntimeUser).mockResolvedValue({ user: { user_id: "owner", display_name: null }, response: null } as never); const sdk = { + getTokens: vi.fn(() => ({ accessToken: "synthetic-test", refreshToken: null })), listRuntimeSessions: vi.fn().mockResolvedValue({ sessions: [session] }), listRuns: vi.fn().mockResolvedValue({ runs: [{ runId: "run", sessionId: "session", parentRunId: null, depth: 0, status: "running" }] }), getSessionConversation: vi.fn().mockResolvedValue(conversation()), diff --git a/web/__tests__/lib/runtime-client-refresh.test.ts b/web/__tests__/lib/runtime-client-refresh.test.ts new file mode 100644 index 0000000000..c8b5d91fa4 --- /dev/null +++ b/web/__tests__/lib/runtime-client-refresh.test.ts @@ -0,0 +1,283 @@ +// @vitest-environment node +vi.mock("next/headers", () => ({ cookies: vi.fn() })); + +import { cookies } from "next/headers"; +import { WebRuntimeClient } from "@/lib/runtime-client/server"; +import { ACCESS_TOKEN_COOKIE, REFRESH_TOKEN_COOKIE, type RuntimeConfig } from "@/lib/runtime-config"; + +beforeEach(() => vi.mocked(cookies).mockReset()); +afterEach(() => vi.unstubAllGlobals()); + +function runtimeConfig(accessToken: string, refreshToken: string): RuntimeConfig { + return { + mode: "live", source: "cookie", apiUrl: "https://example.test", + accessToken, refreshToken, demoMode: false, + hasAccessToken: true, hasRefreshToken: true, message: "", + }; +} + +it("persists refreshed httpOnly cookies before the initiating SDK request completes", async () => { + let release!: () => void; + const available = new Promise(resolve => { release = resolve; }); + const set = vi.fn(); + vi.mocked(cookies).mockImplementation(async () => { + await available; + // This consumer only uses the cookie store's setter. + return { set } as unknown as Awaited>; + }); + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith("/auth/refresh")) { + return Response.json({ access_token: "new-access", refresh_token: "new-refresh" }); + } + return new Headers(init?.headers).get("Authorization") === "Bearer new-access" + ? Response.json({ user_id: "alice" }) : Response.json({}, { status: 401 }); + }); + vi.stubGlobal("fetch", fetchMock); + const runtime = new WebRuntimeClient(runtimeConfig("old-access", "old-refresh")); + const settled = vi.fn(); + const pending = runtime.sdk.getMe(); + void pending.then(settled, settled); + await vi.waitFor(() => expect(cookies).toHaveBeenCalledTimes(1)); + expect(settled).not.toHaveBeenCalled(); + expect(set).not.toHaveBeenCalled(); + await expect(runtime.get("/auth/me", { auth: "required" })).resolves.toEqual({ user_id: "alice" }); + expect(settled).not.toHaveBeenCalled(); + release(); + await expect(pending).resolves.toEqual({ user_id: "alice" }); + const options = { httpOnly: true, sameSite: "lax", path: "/" }; + expect(set.mock.calls).toEqual([ + [REFRESH_TOKEN_COOKIE, "new-refresh", options], + [ACCESS_TOKEN_COOKIE, "new-access", options], + ]); + expect(fetchMock.mock.calls.map(([, init]) => new Headers(init?.headers).get("Authorization"))).toEqual([ + "Bearer old-access", null, "Bearer new-access", "Bearer new-access", + ]); +}); + +it("does not replay a delayed old-token request while cookie persistence fails", async () => { + let release401!: (response: Response) => void; + let releaseCookies!: () => void; + const late401 = new Promise(resolve => { release401 = resolve; }); + const cookiesReady = new Promise(resolve => { releaseCookies = resolve; }); + vi.mocked(cookies).mockImplementation(async () => { + await cookiesReady; + return { set: () => { throw new Error("cookie write unavailable"); } } as unknown as Awaited>; + }); + const expired = () => new Response("expired", { status: 401 }); + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return Response.json({ access_token: 'new', refresh_token: 'rotated' }); + const auth = new Headers(init?.headers).get('Authorization'); + if (url.endsWith('/late') && auth === 'Bearer old') return late401; + return auth === 'Bearer old' ? expired() : new Response('ok'); + }); + vi.stubGlobal('fetch', fetchMock); + const runtime = new WebRuntimeClient(runtimeConfig('old', 'old-refresh')); + const late = runtime.fetchResponse('/late', { auth: 'required' }); + const first = runtime.fetchResponse('/first', { auth: 'required' }); + await vi.waitFor(() => expect(cookies).toHaveBeenCalledTimes(1)); + release401(expired()); + expect((await late).status).toBe(401); + releaseCookies(); + expect((await first).status).toBe(401); + expect(fetchMock.mock.calls.filter(([url]) => url.endsWith('/late'))).toHaveLength(1); + expect(fetchMock.mock.calls.filter(([url]) => url.endsWith('/first'))).toHaveLength(1); +}); + +it("can recover when writing the access cookie fails after the rotated refresh cookie", async () => { + const jar = new Map([[ACCESS_TOKEN_COOKIE, 'old'], [REFRESH_TOKEN_COOKIE, 'old-refresh']]); + let failAccess = true; + vi.mocked(cookies).mockResolvedValue({ + set(name: string, value: string) { + if (name === ACCESS_TOKEN_COOKIE && failAccess) { + failAccess = false; + throw new Error('access cookie write failed'); + } + jar.set(name, value); + }, + } as unknown as Awaited>); + let validRefresh = 'old-refresh'; + let rotations = 0; + const seenRefresh: string[] = []; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + const refresh = JSON.parse(String(init?.body)).refresh_token; + seenRefresh.push(refresh); + if (refresh !== validRefresh) return new Response('expired', { status: 401 }); + rotations++; + validRefresh = `refresh-${rotations}`; + return Response.json({ access_token: `access-${rotations}`, refresh_token: validRefresh }); + } + return new Headers(init?.headers).get('Authorization') === `Bearer access-${rotations}` + ? Response.json({ user_id: 'alice' }) : new Response('expired', { status: 401 }); + })); + const first = new WebRuntimeClient(runtimeConfig('old', 'old-refresh')); + expect((await first.fetchResponse('/resource', { auth: 'required' })).status).toBe(401); + expect(jar.get(ACCESS_TOKEN_COOKIE)).toBe('old'); + expect(jar.get(REFRESH_TOKEN_COOKIE)).toBe('refresh-1'); + const retry = new WebRuntimeClient(runtimeConfig(jar.get(ACCESS_TOKEN_COOKIE)!, jar.get(REFRESH_TOKEN_COOKIE)!)); + await expect(retry.get('/resource', { auth: 'required' })).resolves.toEqual({ user_id: 'alice' }); + expect(seenRefresh).toEqual(['old-refresh', 'refresh-1']); +}); + +it.each([ + { name: "SDK notification", sdkRefresh: true }, + { name: "Web transport", sdkRefresh: false }, +])("does not rotate during an outstanding $name notification and recovers afterward", async ({ sdkRefresh }) => { + let release!: () => void; + const available = new Promise(resolve => { release = resolve; }); + const jar = new Map(); + const set = vi.fn((name: string, value: string) => { jar.set(name, value); }); + vi.mocked(cookies).mockImplementation(async () => { + if (vi.mocked(cookies).mock.calls.length === 1) await available; + return { set } as unknown as Awaited>; + }); + + const issued = new Set(["access-0"]); + const expired = new Set(["access-0"]); + const refreshRequests: string[] = []; + let rotation = 0; + let validRefresh = "refresh-0"; + const unauthorized = () => Response.json({ detail: "Unauthorized" }, { status: 401 }); + vi.stubGlobal("fetch", vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith("/auth/refresh")) { + const token = JSON.parse(String(init?.body)).refresh_token; + refreshRequests.push(token); + if (token !== validRefresh) return unauthorized(); + rotation++; + validRefresh = `refresh-${rotation}`; + const access = `access-${rotation}`; + issued.add(access); + return Response.json({ access_token: access, refresh_token: validRefresh }); + } + const token = new Headers(init?.headers).get("Authorization")?.replace("Bearer ", ""); + return token && issued.has(token) && !expired.has(token) + ? Response.json({ user_id: "alice" }) : unauthorized(); + })); + + const runtime = new WebRuntimeClient(runtimeConfig("access-0", "refresh-0")); + const first = Promise.allSettled([runtime.sdk.getMe()]); + await vi.waitFor(() => expect(cookies).toHaveBeenCalledTimes(1)); + expired.add("access-1"); + if (sdkRefresh) { + await expect(runtime.sdk.getMe()).rejects.toMatchObject({ status: 401 }); + } else { + await expect(runtime.get("/auth/me", { auth: "required" })).rejects.toMatchObject({ status: 401 }); + } + release(); + expect(await first).toMatchObject([{ status: "rejected", reason: { status: 401 } }]); + + expect(set.mock.calls.map(([name, value]) => [name, value])).toEqual([ + [REFRESH_TOKEN_COOKIE, "refresh-1"], + [ACCESS_TOKEN_COOKIE, "access-1"], + ]); + const next = new WebRuntimeClient(runtimeConfig(jar.get(ACCESS_TOKEN_COOKIE)!, jar.get(REFRESH_TOKEN_COOKIE)!)); + await expect(next.sdk.getMe()).resolves.toEqual({ user_id: "alice" }); + expect(jar.get(REFRESH_TOKEN_COOKIE)).toBe(validRefresh); + expect(refreshRequests).toEqual(["refresh-0", "refresh-1"]); +}); + +it("shares one in-flight refresh between SDK and Web requests", async () => { + let release!: () => void; + const refresh = new Promise(resolve => { release = resolve; }); + const set = vi.fn(); + vi.mocked(cookies).mockResolvedValue({ set } as unknown as Awaited>); + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith("/auth/refresh")) { + await refresh; + return Response.json({ access_token: "new", refresh_token: "rotated" }); + } + return new Headers(init?.headers).get("Authorization") === "Bearer new" + ? Response.json({ user_id: "alice" }) : new Response("expired", { status: 401 }); + }); + vi.stubGlobal("fetch", fetchMock); + const runtime = new WebRuntimeClient(runtimeConfig("old", "refresh")); + const pending = Promise.all([runtime.sdk.getMe(), runtime.get("/auth/me", { auth: "required" })]); + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(3)); + release(); + expect(await pending).toEqual([{ user_id: "alice" }, { user_id: "alice" }]); + expect(fetchMock.mock.calls.filter(([url]) => url.endsWith("/auth/refresh"))).toHaveLength(1); + expect(cookies).toHaveBeenCalledTimes(1); + expect(set).toHaveBeenCalledTimes(2); +}); + +it("keeps Web raw responses, body, cancellation and origin routing intact", async () => { + const fetchMock = vi.fn(async (_url: string, _init?: RequestInit) => new Response("data: event\n\n", { status: 403 })); + vi.stubGlobal("fetch", fetchMock); + const runtime = new WebRuntimeClient({ ...runtimeConfig("access", "refresh"), apiUrl: "https://example.test/ignored?query=ignored" }); + const body = new FormData(); + body.set("upload", "contents"); + const signal = new AbortController().signal; + const response = await runtime.fetchResponse("/stream", { method: "POST", body, signal, headers: { Accept: "text/event-stream" } }); + expect(response.bodyUsed).toBe(false); + expect(await response.text()).toBe("data: event\n\n"); + const [url, init] = fetchMock.mock.calls[0]; + expect(url).toBe("https://example.test/stream"); + expect(init).toMatchObject({ method: "POST", body, signal, cache: "no-store" }); + expect(new Headers(init?.headers).get("Accept")).toBe("text/event-stream"); + expect(new Headers(init?.headers).has("Content-Type")).toBe(false); + await expect(runtime.get("/error", { operation: "read item" })).rejects.toMatchObject({ status: 403, operation: "read item" }); + await expect(runtime.sdk.getMe()).rejects.toMatchObject({ status: 403 }); + expect(fetchMock.mock.calls[2][0]).toBe("https://example.test/auth/me"); +}); + +it("uses current SDK credentials for Web auth and keeps old requests fenced on identity replacement", async () => { + let release!: () => void; + const refresh = new Promise(resolve => { release = resolve; }); + const set = vi.fn(); + vi.mocked(cookies).mockResolvedValue({ set } as unknown as Awaited>); + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith("/auth/refresh")) { + await refresh; + return Response.json({ access_token: "stale", refresh_token: "stale-refresh" }); + } + if (url.endsWith("/auth/logout")) return new Response(null, { status: 204 }); + return new Headers(init?.headers).get("Authorization") === "Bearer bob" + ? Response.json({ user_id: "bob" }) : new Response("expired", { status: 401 }); + }); + vi.stubGlobal("fetch", fetchMock); + const runtime = new WebRuntimeClient(runtimeConfig("alice", "refresh")); + const old = Promise.allSettled([runtime.get("/auth/me", { auth: "required" })]); + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); + runtime.sdk.setTokens("bob", "bob-refresh"); + release(); + expect(await old).toMatchObject([{ status: "rejected", reason: { status: 401 } }]); + await expect(runtime.get("/auth/me", { auth: "required" })).resolves.toEqual({ user_id: "bob" }); + expect(set).not.toHaveBeenCalled(); + await runtime.sdk.logout(); + const count = fetchMock.mock.calls.length; + await expect(runtime.get("/auth/me", { auth: "required" })).rejects.toMatchObject({ status: 401 }); + expect(fetchMock).toHaveBeenCalledTimes(count); +}); + +it("does not authenticate or refresh anonymous Web calls and preserves explicit overrides", async () => { + const fetchMock = vi.fn(async (_url: string, _init?: RequestInit) => new Response("unauthorized", { status: 401 })); + vi.stubGlobal("fetch", fetchMock); + const runtime = new WebRuntimeClient(runtimeConfig("access", "refresh")); + await runtime.fetchResponse("/auth/login", { auth: "none", method: "POST", json: { username: "alice" } }); + await runtime.fetchResponse("/auth/login", { auth: "none", headers: { Authorization: "Bearer explicit", "X-User-Id": "hint" } }); + expect(fetchMock).toHaveBeenCalledTimes(2); + const headers = new Headers(fetchMock.mock.calls[0][1]?.headers); + expect(headers.has("Authorization")).toBe(false); + expect(headers.has("X-User-Id")).toBe(false); + expect(headers.get("Content-Type")).toBe("application/json"); + expect(fetchMock.mock.calls[0][1]?.body).toBe('{"username":"alice"}'); + expect(new Headers(fetchMock.mock.calls[1][1]?.headers).get("Authorization")).toBe("Bearer explicit"); + expect(new Headers(fetchMock.mock.calls[1][1]?.headers).get("X-User-Id")).toBe("hint"); +}); + +it("only retries a raw mutation when its caller opts in", async () => { + vi.mocked(cookies).mockResolvedValue({ set: vi.fn() } as unknown as Awaited>); + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return Response.json({ access_token: 'new', refresh_token: 'rotated' }); + return new Headers(init?.headers).get('Authorization') === 'Bearer new' + ? new Response('ok') : new Response('expired', { status: 401 }); + }); + vi.stubGlobal('fetch', fetchMock); + const runtime = new WebRuntimeClient(runtimeConfig('old', 'refresh')); + expect((await runtime.fetchResponse('/write', { method: 'POST', json: { value: 1 } })).status).toBe(401); + expect(fetchMock).toHaveBeenCalledTimes(1); + await expect(runtime.post('/write', { value: 1 }, { retryUnauthorized: false })).rejects.toMatchObject({ status: 401 }); + expect(fetchMock).toHaveBeenCalledTimes(2); + expect((await runtime.fetchResponse('/write', { method: 'POST', json: { value: 1 }, retryUnauthorized: true })).status).toBe(200); + expect(fetchMock).toHaveBeenCalledTimes(5); +}); diff --git a/web/__tests__/lib/web-store-model-cache.test.ts b/web/__tests__/lib/web-store-model-cache.test.ts index 32fed15bfc..906be0aa66 100644 --- a/web/__tests__/lib/web-store-model-cache.test.ts +++ b/web/__tests__/lib/web-store-model-cache.test.ts @@ -30,15 +30,16 @@ function makeRuntime(overrides: Partial<{ accessToken: string | null }> = {}) { tokenExpiresAtMs: null as number | null, }; + const getTokens = vi.fn((): { accessToken: string | null; refreshToken: string | null } => ({ accessToken: mockConfig.accessToken, refreshToken: mockConfig.refreshToken })); MockedWRC.mockImplementation(function () { return { config: mockConfig, - sdk: { listModels: mockListModels }, + sdk: { listModels: mockListModels, getTokens }, }; }); const client = new WebRuntimeClient(mockConfig as any); - return { client, mockListModels }; + return { client, mockListModels, getTokens }; } describe("resolveModelOfferingSelection", () => { @@ -82,6 +83,13 @@ describe("resolveModelOfferingSelection", () => { expect(mockListModels).not.toHaveBeenCalled(); }); + it("checks live credentials after logout rather than initialization config", async () => { + const { client, getTokens } = makeRuntime(); + getTokens.mockReturnValue({ accessToken: null, refreshToken: null }); + await expect(resolveModelOfferingSelection(client, "offering")) + .rejects.toMatchObject({ code: "authentication_required" }); + }); + it("rejects unauthenticated selection without querying the catalog", async () => { const { client, mockListModels } = makeRuntime({ accessToken: null }); diff --git a/web/app/api/chats/[chatId]/stream/route.ts b/web/app/api/chats/[chatId]/stream/route.ts index e47cc52be8..e76ab11839 100644 --- a/web/app/api/chats/[chatId]/stream/route.ts +++ b/web/app/api/chats/[chatId]/stream/route.ts @@ -702,6 +702,7 @@ export async function POST( return runtime.fetchResponse(PATH_CHAT_STREAM, { method: "POST", auth: "required", + retryUnauthorized: true, operation: "stream web chat turn", signal: backendAbortController.signal, json: { diff --git a/web/app/api/works/[workId]/branches/[branchId]/turns/route.ts b/web/app/api/works/[workId]/branches/[branchId]/turns/route.ts index 2eb600c952..4c28cab38e 100644 --- a/web/app/api/works/[workId]/branches/[branchId]/turns/route.ts +++ b/web/app/api/works/[workId]/branches/[branchId]/turns/route.ts @@ -148,6 +148,7 @@ export async function POST( const backend = await runtime.fetchResponse(path, { method: "POST", auth: "required", + retryUnauthorized: true, body: JSON.stringify(input), signal: request.signal, headers: { diff --git a/web/lib/api/web-store.ts b/web/lib/api/web-store.ts index 2fcf9c8e91..f326381ade 100644 --- a/web/lib/api/web-store.ts +++ b/web/lib/api/web-store.ts @@ -2347,7 +2347,7 @@ export async function resolveModelOfferingSelection( runtime instanceof WebRuntimeClient ? runtime : new WebRuntimeClient(runtime); - const accessToken = client.config.accessToken; + const { accessToken } = client.sdk.getTokens(); if (!accessToken) { throw new ModelOfferingSelectionError( "authentication_required", diff --git a/web/lib/runtime-client/server.ts b/web/lib/runtime-client/server.ts index 7a0923d9f8..a771a39f17 100644 --- a/web/lib/runtime-client/server.ts +++ b/web/lib/runtime-client/server.ts @@ -1,9 +1,7 @@ import { AstraClient, - PATH_AUTH_REFRESH, extractJwtSubject, headersInitToRecord, - joinApiPath, methodCanHaveJson, } from "@astra/sdk"; import { @@ -21,6 +19,7 @@ export type RuntimeAuthMode = "required" | "optional" | "none"; export type RuntimeRequestInit = Omit & { auth?: RuntimeAuthMode; + retryUnauthorized?: boolean; json?: unknown; operation?: string; body?: BodyInit | null; @@ -34,8 +33,6 @@ type RuntimeClientOptions = { export class WebRuntimeClient { readonly config: RuntimeConfig; readonly sdk: AstraClient; - private accessToken?: string; - private refreshToken?: string; constructor(config: RuntimeConfig) { if (config.mode !== "live" || !config.apiUrl) { @@ -48,25 +45,16 @@ export class WebRuntimeClient { } this.config = config; - this.accessToken = config.accessToken; - this.refreshToken = config.refreshToken; this.sdk = new AstraClient({ - baseUrl: config.apiUrl, + baseUrl: new URL("/", config.apiUrl).toString(), accessToken: config.accessToken, refreshToken: config.refreshToken, - headers: this.baseAuthHeaders("optional"), onTokenRefresh: async (tokens) => { - this.accessToken = tokens.accessToken; - this.refreshToken = tokens.refreshToken; await this.persistTokens(tokens.accessToken, tokens.refreshToken); }, }); } - get apiUrl(): string { - return this.config.apiUrl ?? DEFAULT_API_URL; - } - async get(path: string, init?: RuntimeRequestInit): Promise { return this.request(path, { ...init, method: "GET" }); } @@ -92,7 +80,7 @@ export class WebRuntimeClient { } async request(path: string, init: RuntimeRequestInit = {}): Promise { - const response = await this.fetchResponse(path, init); + const response = await this.fetchResponse(path, { ...init, retryUnauthorized: init.retryUnauthorized ?? true }); if (!response.ok) { throw new RuntimeClientError({ operation: init.operation ?? `${init.method ?? "GET"} ${path}`, @@ -130,24 +118,11 @@ export class WebRuntimeClient { ): Promise { const operation = init.operation ?? `${init.method ?? "GET"} ${path}`; const auth = init.auth ?? "optional"; - const url = this.url(path); - const request = this.toRequestInit(init, auth, operation, path); - - let response = await fetch(url, request); - if (response.status === 401 && auth !== "none" && this.refreshToken) { - const refreshed = await this.refreshTokens(); - if (refreshed) { - response = await fetch( - url, - this.toRequestInit(init, auth, operation, path), - ); - } - } - return response; - } - - url(path: string): string { - return new URL(joinApiPath(undefined, path), this.apiUrl).toString(); + return this.sdk.fetchResponse( + path, + this.toRequestInit(init, auth, operation, path), + { auth: auth === "none" ? "none" : "optional", retryUnauthorized: init.retryUnauthorized }, + ); } private toRequestInit( @@ -167,7 +142,7 @@ export class WebRuntimeClient { init.headers, ); - if (auth === "required" && !this.accessToken) { + if (auth === "required" && !this.sdk.getTokens().accessToken) { throw new RuntimeClientError({ operation, path, @@ -176,7 +151,7 @@ export class WebRuntimeClient { }); } - const { auth: _auth, json: _json, operation: _operation, ...rest } = init; + const { auth: _auth, json: _json, operation: _operation, retryUnauthorized: _retryUnauthorized, ...rest } = init; return { ...rest, method, @@ -194,9 +169,9 @@ export class WebRuntimeClient { if (request?.hasJsonBody && methodCanHaveJson(request.method)) { headers["Content-Type"] = "application/json"; } - if (auth !== "none" && this.accessToken) { - headers.Authorization = `Bearer ${this.accessToken}`; - const userId = extractJwtSubject(this.accessToken); + const { accessToken } = this.sdk.getTokens(); + if (auth !== "none" && accessToken) { + const userId = extractJwtSubject(accessToken); if (userId) { headers["X-User-Id"] = userId; } @@ -204,51 +179,23 @@ export class WebRuntimeClient { return headers; } - private async refreshTokens(): Promise { - if (!this.refreshToken) { - return false; - } - - const response = await fetch(this.url(PATH_AUTH_REFRESH), { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ refresh_token: this.refreshToken }), - cache: "no-store", - }); - if (!response.ok) { - return false; - } - - const payload = (await response.json()) as { - access_token?: unknown; - refresh_token?: unknown; - }; - if ( - typeof payload.access_token !== "string" || - typeof payload.refresh_token !== "string" - ) { - return false; - } - - this.accessToken = payload.access_token; - this.refreshToken = payload.refresh_token; - this.sdk.setTokens(payload.access_token, payload.refresh_token); - await this.persistTokens(payload.access_token, payload.refresh_token); - return true; - } - private async persistTokens( accessToken: string, refreshToken: string, ): Promise { const { cookies } = await import("next/headers"); const cookieStore = await cookies(); - cookieStore.set(ACCESS_TOKEN_COOKIE, accessToken, { + // A newer refresh may have replaced this pair while cookie acquisition waited. + const current = this.sdk.getTokens(); + if (current.accessToken !== accessToken || current.refreshToken !== refreshToken) { + return; + } + cookieStore.set(REFRESH_TOKEN_COOKIE, refreshToken, { httpOnly: true, sameSite: "lax", path: "/", }); - cookieStore.set(REFRESH_TOKEN_COOKIE, refreshToken, { + cookieStore.set(ACCESS_TOKEN_COOKIE, accessToken, { httpOnly: true, sameSite: "lax", path: "/",