Skip to content

Commit a3129a1

Browse files
gpealgpeal
committed
Use key
1 parent 5ec53ce commit a3129a1

File tree

2 files changed

+50
-44
lines changed

2 files changed

+50
-44
lines changed

‎codex-rs/cli/src/mcp_cmd.rs‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -246,11 +246,17 @@ fn run_logout(config_overrides: &CliConfigOverrides, logout_args: LogoutArgs) ->
246246

247247
let LogoutArgs { name } = logout_args;
248248

249-
if !config.mcp_servers.contains_key(&name) {
250-
println!("No MCP server named '{name}' found in configuration.");
251-
}
249+
let server = config
250+
.mcp_servers
251+
.get(&name)
252+
.ok_or_else(|| anyhow!("No MCP server named '{name}' found in configuration."))?;
253+
254+
let url = match &server.transport {
255+
McpServerTransportConfig::StreamableHttp { url, .. } => url.clone(),
256+
_ => bail!("OAuth logout is only supported for streamable_http transports."),
257+
};
252258

253-
match delete_oauth_tokens(&name) {
259+
match delete_oauth_tokens(&name, &url) {
254260
Ok(true) => println!("Removed OAuth credentials for '{name}'."),
255261
Ok(false) => println!("No OAuth credentials stored for '{name}'."),
256262
Err(err) => return Err(anyhow!("failed to delete OAuth credentials: {err}")),

‎codex-rs/rmcp-client/src/oauth.rs‎

Lines changed: 40 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,8 @@ fn load_oauth_tokens_with_store<C: CredentialStore>(
139139
server_name: &str,
140140
url: &str,
141141
) -> Result<Option<StoredOAuthTokens>> {
142-
match store.load(KEYRING_SERVICE, server_name) {
142+
let key = compute_store_key(server_name, url)?;
143+
match store.load(KEYRING_SERVICE, &key) {
143144
Ok(Some(serialized)) => {
144145
let tokens: StoredOAuthTokens = serde_json::from_str(&serialized)
145146
.context("failed to deserialize OAuth tokens from keyring")?;
@@ -167,14 +168,19 @@ fn save_oauth_tokens_with_store<C: CredentialStore>(
167168
) -> Result<()> {
168169
let serialized = serde_json::to_string(tokens).context("failed to serialize OAuth tokens")?;
169170

170-
match store.save(KEYRING_SERVICE, server_name, &serialized) {
171+
let key = compute_store_key(server_name, &tokens.url)?;
172+
// DO NOT SUBMIT
173+
println!("[DEBUG] saving OAuth tokens to keyring: {key}: {serialized}");
174+
match store.save(KEYRING_SERVICE, &key, &serialized) {
171175
Ok(()) => {
172-
if let Err(error) = delete_oauth_tokens_from_file(server_name, Some(&tokens.url)) {
176+
println!("saved OAuth tokens to keyring: {serialized}");
177+
if let Err(error) = delete_oauth_tokens_from_file(&key) {
173178
warn!("failed to remove OAuth tokens from fallback storage: {error:?}");
174179
}
175180
Ok(())
176181
}
177182
Err(error) => {
183+
println!("failed to save OAuth tokens to keyring: {error:?}");
178184
let message = error.message();
179185
warn!("failed to write OAuth tokens to keyring: {message}");
180186
save_oauth_tokens_to_file(tokens)
@@ -183,16 +189,18 @@ fn save_oauth_tokens_with_store<C: CredentialStore>(
183189
}
184190
}
185191

186-
pub fn delete_oauth_tokens(server_name: &str) -> Result<bool> {
192+
pub fn delete_oauth_tokens(server_name: &str, url: &str) -> Result<bool> {
187193
let store = KeyringCredentialStore;
188-
delete_oauth_tokens_with_store(&store, server_name)
194+
delete_oauth_tokens_with_store(&store, server_name, url)
189195
}
190196

191197
fn delete_oauth_tokens_with_store<C: CredentialStore>(
192198
store: &C,
193199
server_name: &str,
200+
url: &str,
194201
) -> Result<bool> {
195-
let keyring_removed = match store.delete(KEYRING_SERVICE, server_name) {
202+
let key = compute_store_key(server_name, url)?;
203+
let keyring_removed = match store.delete(KEYRING_SERVICE, &key) {
196204
Ok(removed) => removed,
197205
Err(error) => {
198206
let message = error.message();
@@ -201,7 +209,7 @@ fn delete_oauth_tokens_with_store<C: CredentialStore>(
201209
}
202210
};
203211

204-
let file_removed = delete_oauth_tokens_from_file(server_name, None)?;
212+
let file_removed = delete_oauth_tokens_from_file(&key)?;
205213
Ok(keyring_removed || file_removed)
206214
}
207215

@@ -260,7 +268,8 @@ impl OAuthRuntime {
260268
None => {
261269
let mut last_serialized = self.inner.last_credentials.lock().await;
262270
if last_serialized.take().is_some()
263-
&& let Err(error) = delete_oauth_tokens(&self.inner.server_name)
271+
&& let Err(error) =
272+
delete_oauth_tokens(&self.inner.server_name, &self.inner.url)
264273
{
265274
warn!(
266275
"failed to remove OAuth tokens for server {}: {error}",
@@ -342,8 +351,8 @@ fn load_oauth_tokens_from_file(server_name: &str, url: &str) -> Result<Option<St
342351
}
343352

344353
fn save_oauth_tokens_to_file(tokens: &StoredOAuthTokens) -> Result<()> {
345-
let mut store = read_fallback_file()?.unwrap_or_default();
346354
let key = compute_store_key(&tokens.server_name, &tokens.url)?;
355+
let mut store = read_fallback_file()?.unwrap_or_default();
347356

348357
let token_response = &tokens.token_response.0;
349358
let refresh_token = token_response
@@ -367,20 +376,13 @@ fn save_oauth_tokens_to_file(tokens: &StoredOAuthTokens) -> Result<()> {
367376
write_fallback_file(&store)
368377
}
369378

370-
fn delete_oauth_tokens_from_file(server_name: &str, url: Option<&str>) -> Result<bool> {
379+
fn delete_oauth_tokens_from_file(key: &str) -> Result<bool> {
371380
let mut store = match read_fallback_file()? {
372381
Some(store) => store,
373382
None => return Ok(false),
374383
};
375384

376-
let removed = if let Some(target_url) = url {
377-
let key = compute_store_key(server_name, target_url)?;
378-
store.remove(&key).is_some()
379-
} else {
380-
let original_len = store.len();
381-
store.retain(|_, entry| entry.server_name != server_name);
382-
store.len() != original_len
383-
};
385+
let removed = store.remove(key).is_some();
384386

385387
if removed {
386388
write_fallback_file(&store)?;
@@ -642,7 +644,8 @@ mod tests {
642644
let tokens = sample_tokens();
643645
let expected = tokens.clone();
644646
let serialized = serde_json::to_string(&tokens)?;
645-
store.save(KEYRING_SERVICE, &tokens.server_name, &serialized)?;
647+
let key = super::compute_store_key(&tokens.server_name, &tokens.url)?;
648+
store.save(KEYRING_SERVICE, &key, &serialized)?;
646649

647650
let loaded = super::load_oauth_tokens_with_store(&store, &tokens.server_name, &tokens.url)?;
648651
assert_eq!(loaded, Some(expected));
@@ -670,10 +673,8 @@ mod tests {
670673
let store = MockCredentialStore::default();
671674
let tokens = sample_tokens();
672675
let expected = tokens.clone();
673-
store.set_error(
674-
&tokens.server_name,
675-
KeyringError::Invalid("error".into(), "load".into()),
676-
);
676+
let key = super::compute_store_key(&tokens.server_name, &tokens.url)?;
677+
store.set_error(&key, KeyringError::Invalid("error".into(), "load".into()));
677678

678679
super::save_oauth_tokens_to_file(&tokens)?;
679680

@@ -688,16 +689,15 @@ mod tests {
688689
let _env = TempCodexHome::new();
689690
let store = MockCredentialStore::default();
690691
let tokens = sample_tokens();
692+
let key = super::compute_store_key(&tokens.server_name, &tokens.url)?;
691693

692694
super::save_oauth_tokens_to_file(&tokens)?;
693695

694696
super::save_oauth_tokens_with_store(&store, &tokens.server_name, &tokens)?;
695697

696698
let fallback_path = super::fallback_file_path()?;
697699
assert!(!fallback_path.exists(), "fallback file should be removed");
698-
let stored = store
699-
.saved_value(&tokens.server_name)
700-
.expect("value saved to keyring");
700+
let stored = store.saved_value(&key).expect("value saved to keyring");
701701
assert_eq!(serde_json::from_str::<StoredOAuthTokens>(&stored)?, tokens);
702702
Ok(())
703703
}
@@ -707,10 +707,8 @@ mod tests {
707707
let _env = TempCodexHome::new();
708708
let store = MockCredentialStore::default();
709709
let tokens = sample_tokens();
710-
store.set_error(
711-
&tokens.server_name,
712-
KeyringError::Invalid("error".into(), "save".into()),
713-
);
710+
let key = super::compute_store_key(&tokens.server_name, &tokens.url)?;
711+
store.set_error(&key, KeyringError::Invalid("error".into(), "save".into()));
714712

715713
super::save_oauth_tokens_with_store(&store, &tokens.server_name, &tokens)?;
716714

@@ -726,7 +724,7 @@ mod tests {
726724
entry.access_token,
727725
tokens.token_response.0.access_token().secret().as_str()
728726
);
729-
assert!(store.saved_value(&tokens.server_name).is_none());
727+
assert!(store.saved_value(&key).is_none());
730728
Ok(())
731729
}
732730

@@ -736,30 +734,32 @@ mod tests {
736734
let store = MockCredentialStore::default();
737735
let tokens = sample_tokens();
738736
let serialized = serde_json::to_string(&tokens)?;
739-
store.save(KEYRING_SERVICE, &tokens.server_name, &serialized)?;
737+
let key = super::compute_store_key(&tokens.server_name, &tokens.url)?;
738+
store.save(KEYRING_SERVICE, &key, &serialized)?;
740739
super::save_oauth_tokens_to_file(&tokens)?;
741740

742-
let removed = super::delete_oauth_tokens_with_store(&store, &tokens.server_name)?;
741+
let removed =
742+
super::delete_oauth_tokens_with_store(&store, &tokens.server_name, &tokens.url)?;
743743
assert!(removed);
744-
assert!(!store.contains(&tokens.server_name));
744+
assert!(!store.contains(&key));
745745
assert!(!super::fallback_file_path()?.exists());
746746
Ok(())
747747
}
748748

749749
#[test]
750-
fn delete_oauth_tokens_propagates_keyring_errors() {
750+
fn delete_oauth_tokens_propagates_keyring_errors() -> Result<()> {
751751
let _env = TempCodexHome::new();
752752
let store = MockCredentialStore::default();
753753
let tokens = sample_tokens();
754-
store.set_error(
755-
&tokens.server_name,
756-
KeyringError::Invalid("error".into(), "delete".into()),
757-
);
754+
let key = super::compute_store_key(&tokens.server_name, &tokens.url)?;
755+
store.set_error(&key, KeyringError::Invalid("error".into(), "delete".into()));
758756
super::save_oauth_tokens_to_file(&tokens).unwrap();
759757

760-
let result = super::delete_oauth_tokens_with_store(&store, &tokens.server_name);
758+
let result =
759+
super::delete_oauth_tokens_with_store(&store, &tokens.server_name, &tokens.url);
761760
assert!(result.is_err());
762761
assert!(super::fallback_file_path().unwrap().exists());
762+
Ok(())
763763
}
764764

765765
fn assert_tokens_match_without_expiry(

0 commit comments

Comments
 (0)