rmcp-client to use json store api

This commit is contained in:
mikhail-oai
2026-03-24 09:55:12 -04:00
parent 4b01a5ee4e
commit 60145f0f09

View File

@@ -45,6 +45,9 @@ use tracing::warn;
use codex_keyring_store::DefaultKeyringStore;
use codex_keyring_store::KeyringStore;
use codex_keyring_store::delete_json_from_keyring;
use codex_keyring_store::load_json_from_keyring;
use codex_keyring_store::save_json_to_keyring;
use rmcp::transport::auth::AuthorizationManager;
use tokio::sync::Mutex;
@@ -155,16 +158,15 @@ fn load_oauth_tokens_from_keyring<K: KeyringStore>(
url: &str,
) -> Result<Option<StoredOAuthTokens>> {
let key = compute_store_key(server_name, url)?;
match keyring_store.load(KEYRING_SERVICE, &key) {
Ok(Some(serialized)) => {
let mut tokens: StoredOAuthTokens = serde_json::from_str(&serialized)
.context("failed to deserialize OAuth tokens from keyring")?;
refresh_expires_in_from_timestamp(&mut tokens);
Ok(Some(tokens))
}
Ok(None) => Ok(None),
Err(error) => Err(Error::new(error.into_error())),
}
let Some(value) = load_json_from_keyring(keyring_store, KEYRING_SERVICE, &key)
.map_err(|err| Error::msg(err.to_string()))?
else {
return Ok(None);
};
let mut tokens: StoredOAuthTokens =
serde_json::from_value(value).context("failed to deserialize OAuth tokens from keyring")?;
refresh_expires_in_from_timestamp(&mut tokens);
Ok(Some(tokens))
}
pub fn save_oauth_tokens(
@@ -191,10 +193,9 @@ fn save_oauth_tokens_with_keyring<K: KeyringStore>(
server_name: &str,
tokens: &StoredOAuthTokens,
) -> Result<()> {
let serialized = serde_json::to_string(tokens).context("failed to serialize OAuth tokens")?;
let value = serde_json::to_value(tokens).context("failed to serialize OAuth tokens")?;
let key = compute_store_key(server_name, &tokens.url)?;
match keyring_store.save(KEYRING_SERVICE, &key, &serialized) {
match save_json_to_keyring(keyring_store, KEYRING_SERVICE, &key, &value) {
Ok(()) => {
if let Err(error) = delete_oauth_tokens_from_file(&key) {
warn!("failed to remove OAuth tokens from fallback storage: {error:?}");
@@ -202,12 +203,9 @@ fn save_oauth_tokens_with_keyring<K: KeyringStore>(
Ok(())
}
Err(error) => {
let message = format!(
"failed to write OAuth tokens to keyring: {}",
error.message()
);
let message = format!("failed to write OAuth tokens to keyring: {error}");
warn!("{message}");
Err(Error::new(error.into_error()).context(message))
Err(Error::msg(message))
}
}
}
@@ -244,22 +242,20 @@ fn delete_oauth_tokens_from_keyring_and_file<K: KeyringStore>(
url: &str,
) -> Result<bool> {
let key = compute_store_key(server_name, url)?;
let keyring_result = keyring_store.delete(KEYRING_SERVICE, &key);
let keyring_removed = match keyring_result {
let keyring_removed = match delete_json_from_keyring(keyring_store, KEYRING_SERVICE, &key) {
Ok(removed) => removed,
Err(error) => {
let message = error.message();
let message = error.to_string();
warn!("failed to delete OAuth tokens from keyring: {message}");
match store_mode {
OAuthCredentialsStoreMode::Auto | OAuthCredentialsStoreMode::Keyring => {
return Err(error.into_error())
return Err(Error::msg(message))
.context("failed to delete OAuth tokens from keyring");
}
OAuthCredentialsStoreMode::File => false,
}
}
};
let file_removed = delete_oauth_tokens_from_file(&key)?;
Ok(keyring_removed || file_removed)
}
@@ -604,6 +600,9 @@ fn sha_256_prefix(value: &Value) -> Result<String> {
mod tests {
use super::*;
use anyhow::Result;
use codex_keyring_store::CredentialStoreError;
use codex_keyring_store::load_json_from_keyring;
use codex_keyring_store::save_json_to_keyring;
use keyring::Error as KeyringError;
use pretty_assertions::assert_eq;
use std::sync::Mutex;
@@ -614,6 +613,101 @@ mod tests {
use codex_keyring_store::tests::MockKeyringStore;
#[derive(Clone, Debug)]
struct KeyringStoreWithError {
inner: MockKeyringStore,
fail_delete: bool,
fail_load_secret: bool,
fail_save_secret: bool,
}
impl KeyringStoreWithError {
fn fail_delete(inner: MockKeyringStore) -> Self {
Self {
inner,
fail_delete: true,
fail_load_secret: false,
fail_save_secret: false,
}
}
fn fail_load_secret(inner: MockKeyringStore) -> Self {
Self {
inner,
fail_delete: false,
fail_load_secret: true,
fail_save_secret: false,
}
}
fn fail_save_secret(inner: MockKeyringStore) -> Self {
Self {
inner,
fail_delete: false,
fail_load_secret: false,
fail_save_secret: true,
}
}
}
impl KeyringStore for KeyringStoreWithError {
fn load(
&self,
service: &str,
account: &str,
) -> Result<Option<String>, CredentialStoreError> {
self.inner.load(service, account)
}
fn load_secret(
&self,
service: &str,
account: &str,
) -> Result<Option<Vec<u8>>, CredentialStoreError> {
if self.fail_load_secret {
return Err(CredentialStoreError::new(KeyringError::Invalid(
"error".into(),
"load".into(),
)));
}
self.inner.load_secret(service, account)
}
fn save(
&self,
service: &str,
account: &str,
value: &str,
) -> Result<(), CredentialStoreError> {
self.inner.save(service, account, value)
}
fn save_secret(
&self,
service: &str,
account: &str,
value: &[u8],
) -> Result<(), CredentialStoreError> {
if self.fail_save_secret {
return Err(CredentialStoreError::new(KeyringError::Invalid(
"error".into(),
"save".into(),
)));
}
self.inner.save_secret(service, account, value)
}
fn delete(&self, service: &str, account: &str) -> Result<bool, CredentialStoreError> {
if self.fail_delete {
return Err(CredentialStoreError::new(KeyringError::Invalid(
"error".into(),
"delete".into(),
)));
}
self.inner.delete(service, account)
}
}
struct TempCodexHome {
_guard: MutexGuard<'static, ()>,
_dir: tempfile::TempDir,
@@ -651,9 +745,9 @@ mod tests {
let store = MockKeyringStore::default();
let tokens = sample_tokens();
let expected = tokens.clone();
let serialized = serde_json::to_string(&tokens)?;
let key = super::compute_store_key(&tokens.server_name, &tokens.url)?;
store.save(KEYRING_SERVICE, &key, &serialized)?;
let value = serde_json::to_value(&tokens)?;
save_json_to_keyring(&store, KEYRING_SERVICE, &key, &value)?;
let loaded =
super::load_oauth_tokens_from_keyring(&store, &tokens.server_name, &tokens.url)?
@@ -662,6 +756,31 @@ mod tests {
Ok(())
}
#[test]
fn load_oauth_tokens_supports_legacy_single_entry() -> Result<()> {
let _env = TempCodexHome::new();
let store = MockKeyringStore::default();
let tokens = sample_tokens();
let serialized = serde_json::to_string(&tokens)?;
let key = super::compute_store_key(&tokens.server_name, &tokens.url)?;
store.save(KEYRING_SERVICE, &key, &serialized)?;
let loaded =
super::load_oauth_tokens_from_keyring(&store, &tokens.server_name, &tokens.url)?;
#[cfg(not(windows))]
{
let loaded = loaded.expect("tokens should load from keyring");
assert_tokens_match_without_expiry(&loaded, &tokens);
}
#[cfg(windows)]
{
assert!(loaded.is_none());
}
Ok(())
}
#[test]
fn load_oauth_tokens_falls_back_when_missing_in_keyring() -> Result<()> {
let _env = TempCodexHome::new();
@@ -684,11 +803,9 @@ mod tests {
#[test]
fn load_oauth_tokens_falls_back_when_keyring_errors() -> Result<()> {
let _env = TempCodexHome::new();
let store = MockKeyringStore::default();
let store = KeyringStoreWithError::fail_load_secret(MockKeyringStore::default());
let tokens = sample_tokens();
let expected = tokens.clone();
let key = super::compute_store_key(&tokens.server_name, &tokens.url)?;
store.set_error(&key, KeyringError::Invalid("error".into(), "load".into()));
super::save_oauth_tokens_to_file(&tokens)?;
@@ -719,18 +836,29 @@ mod tests {
let fallback_path = super::fallback_file_path()?;
assert!(!fallback_path.exists(), "fallback file should be removed");
let stored = store.saved_value(&key).expect("value saved to keyring");
assert_eq!(serde_json::from_str::<StoredOAuthTokens>(&stored)?, tokens);
#[cfg(windows)]
assert!(
store.saved_secret(&key).is_none(),
"windows should not store the full JSON record under the base key"
);
#[cfg(not(windows))]
assert!(
store.saved_secret(&key).is_some(),
"non-windows should store the full JSON record as one secret"
);
let stored =
super::load_oauth_tokens_from_keyring(&store, &tokens.server_name, &tokens.url)?
.expect("value saved to keyring");
assert_tokens_match_without_expiry(&stored, &tokens);
Ok(())
}
#[test]
fn save_oauth_tokens_writes_fallback_when_keyring_fails() -> Result<()> {
let _env = TempCodexHome::new();
let store = MockKeyringStore::default();
let mock_keyring = MockKeyringStore::default();
let store = KeyringStoreWithError::fail_save_secret(mock_keyring.clone());
let tokens = sample_tokens();
let key = super::compute_store_key(&tokens.server_name, &tokens.url)?;
store.set_error(&key, KeyringError::Invalid("error".into(), "save".into()));
super::save_oauth_tokens_with_keyring_with_fallback_to_file(
&store,
@@ -750,18 +878,22 @@ mod tests {
entry.access_token,
tokens.token_response.0.access_token().secret().as_str()
);
assert!(store.saved_value(&key).is_none());
assert!(mock_keyring.saved_value(&key).is_none());
assert!(
load_json_from_keyring(&mock_keyring, KEYRING_SERVICE, &key)?.is_none(),
"keyring should not point at saved OAuth tokens when save fails"
);
Ok(())
}
#[test]
fn delete_oauth_tokens_removes_all_storage() -> Result<()> {
fn delete_oauth_tokens_removes_active_storage() -> Result<()> {
let _env = TempCodexHome::new();
let store = MockKeyringStore::default();
let tokens = sample_tokens();
let serialized = serde_json::to_string(&tokens)?;
let key = super::compute_store_key(&tokens.server_name, &tokens.url)?;
store.save(KEYRING_SERVICE, &key, &serialized)?;
let value = serde_json::to_value(&tokens)?;
save_json_to_keyring(&store, KEYRING_SERVICE, &key, &value)?;
super::save_oauth_tokens_to_file(&tokens)?;
let removed = super::delete_oauth_tokens_from_keyring_and_file(
@@ -771,7 +903,10 @@ mod tests {
&tokens.url,
)?;
assert!(removed);
assert!(!store.contains(&key));
assert!(
load_json_from_keyring(&store, KEYRING_SERVICE, &key)?.is_none(),
"keyring entry should be removed"
);
assert!(!super::fallback_file_path()?.exists());
Ok(())
}
@@ -781,10 +916,13 @@ mod tests {
let _env = TempCodexHome::new();
let store = MockKeyringStore::default();
let tokens = sample_tokens();
let serialized = serde_json::to_string(&tokens)?;
let key = super::compute_store_key(&tokens.server_name, &tokens.url)?;
store.save(KEYRING_SERVICE, &key, &serialized)?;
assert!(store.contains(&key));
let value = serde_json::to_value(&tokens)?;
save_json_to_keyring(&store, KEYRING_SERVICE, &key, &value)?;
assert!(
super::load_oauth_tokens_from_keyring(&store, &tokens.server_name, &tokens.url)?
.is_some()
);
let removed = super::delete_oauth_tokens_from_keyring_and_file(
&store,
@@ -793,7 +931,14 @@ mod tests {
&tokens.url,
)?;
assert!(removed);
assert!(!store.contains(&key));
assert!(
load_json_from_keyring(&store, KEYRING_SERVICE, &key)?.is_none(),
"keyring entry should be removed"
);
assert!(
super::load_oauth_tokens_from_keyring(&store, &tokens.server_name, &tokens.url)?
.is_none()
);
assert!(!super::fallback_file_path()?.exists());
Ok(())
}
@@ -801,10 +946,11 @@ mod tests {
#[test]
fn delete_oauth_tokens_propagates_keyring_errors() -> Result<()> {
let _env = TempCodexHome::new();
let store = MockKeyringStore::default();
let store = KeyringStoreWithError::fail_delete(MockKeyringStore::default());
let tokens = sample_tokens();
let key = super::compute_store_key(&tokens.server_name, &tokens.url)?;
store.set_error(&key, KeyringError::Invalid("error".into(), "delete".into()));
let value = serde_json::to_value(&tokens)?;
save_json_to_keyring(&store, KEYRING_SERVICE, &key, &value)?;
super::save_oauth_tokens_to_file(&tokens).unwrap();
let result = super::delete_oauth_tokens_from_keyring_and_file(