Add attributed MCP OAuth credential storage telemetry (#49392)

## What changed

- Instrument MCP OAuth loads, saves, deletes, fallback cleanup, and refresh persistence with credential-free storage metrics, distinguishing configured policy from operations on a pinned store.
- Retain the originating client across background tasks, blocking storage operations, and session recovery, including recovery after startup found no credentials.
- Preserve typed keyring and lock error categories for metrics and remove raw error details from storage warnings.

## Testing

Add tests for policy and pinned-store outcomes, best-effort cleanup failures, wrapped lock errors, and successful or failed refresh persistence without duplicate observations. Add recovery tests that verify client attribution after anonymous startup and across threads with different originators.

GitOrigin-RevId: 5ea4528bcdcf500c8ffdf1d9c7008ea52112c726
This commit is contained in:
Celia Chen
2026-09-29 23:01:46 +00:00
committed by copyberry
parent 9212b3eca8
commit 05ea5f757e
25 changed files with 1279 additions and 187 deletions
+1
View File
@@ -4389,6 +4389,7 @@ dependencies = [
"libc",
"memchr",
"oauth2",
"opentelemetry_sdk",
"pretty_assertions",
"rmcp",
"serde",
@@ -7,6 +7,7 @@ use codex_mcp::McpServerSource;
use codex_mcp::ReadResourceRequestParams;
use codex_mcp::ema_auth_scope;
use codex_mcp::resolve_oauth_callback;
use codex_otel::auth_storage::AuthStorageOriginator;
use codex_rmcp_client::EnterpriseOAuthLoginRequest;
use std::future::Future;
@@ -379,12 +380,13 @@ impl McpRequestProcessor {
None => None,
};
let processor = self.clone();
tokio::spawn(async move {
let task = async move {
let result = processor
.list_mcp_server_status_response(params, thread)
.await;
outgoing.send_result(request, result).await;
});
};
tokio::spawn(AuthStorageOriginator::current().scope(task));
Ok(())
}
@@ -598,7 +600,7 @@ impl McpRequestProcessor {
let (_, thread) = self.load_thread(&thread_id).await?;
let request_id = request_id.clone();
tokio::spawn(async move {
let task = async move {
let origin_call_id =
origin_call_id.filter(|_| server == codex_mcp::CODEX_APPS_MCP_SERVER_NAME);
let result = match origin_call_id.as_deref() {
@@ -611,7 +613,8 @@ impl McpRequestProcessor {
};
Self::send_mcp_resource_read_response(outgoing, request_id, result, origin_call_id)
.await;
});
};
tokio::spawn(AuthStorageOriginator::current().scope(task));
return Ok(());
}
@@ -633,7 +636,7 @@ impl McpRequestProcessor {
McpRuntimeContext::new(Arc::clone(&environment_manager), config.cwd.to_path_buf());
let request_id = request_id.clone();
tokio::spawn(async move {
let task = async move {
let result = read_mcp_resource_without_thread(
&mcp_config,
auth.as_ref(),
@@ -649,7 +652,8 @@ impl McpRequestProcessor {
outgoing, request_id, result, /*origin_call_id*/ None,
)
.await;
});
};
tokio::spawn(AuthStorageOriginator::current().scope(task));
Ok(())
}
@@ -687,14 +691,15 @@ impl McpRequestProcessor {
let meta = with_mcp_tool_call_thread_id_meta(params.meta, &thread_id);
let request_id = request_id.clone();
tokio::spawn(async move {
let task = async move {
let result = thread
.call_mcp_tool(&params.server, &params.tool, params.arguments, meta)
.await
.map(McpServerToolCallResponse::from)
.map_err(mcp_operation_error);
outgoing.send_result(request_id, result).await;
});
};
tokio::spawn(AuthStorageOriginator::current().scope(task));
Ok(())
}
}
+13 -8
View File
@@ -65,6 +65,7 @@ use codex_config::McpServerTransportConfig;
use codex_config::McpStartupReadiness;
use codex_diagnostics::Gauge;
use codex_diagnostics::GaugeGuard;
use codex_otel::auth_storage::AuthStorageOriginator;
use codex_protocol::mcp::CallToolResult;
use codex_protocol::mcp::McpServerInfo;
use codex_protocol::protocol::Event;
@@ -772,6 +773,7 @@ impl McpConnectionSet {
(server_name, outcome)
};
let startup = AuthStorageOriginator::current().scope(startup);
if defer_startup {
// Dormant servers must not hold the initial startup summary open.
tokio::spawn(startup);
@@ -893,15 +895,18 @@ impl McpConnectionSet {
return Vec::new();
}
let originator = AuthStorageOriginator::current();
match tokio::task::spawn_blocking(move || {
candidates
.into_iter()
.filter_map(|(server_name, identity, config)| {
identity
.oauth_credentials_changed(&server_name, &config)
.then_some(server_name)
})
.collect()
originator.sync_scope(|| {
candidates
.into_iter()
.filter_map(|(server_name, identity, config)| {
identity
.oauth_credentials_changed(&server_name, &config)
.then_some(server_name)
})
.collect()
})
})
.await
{
+7 -5
View File
@@ -57,6 +57,7 @@ use codex_connectors::ConnectorRuntimeContext;
use codex_connectors::ConnectorRuntimeFetchSource;
use codex_exec_server::Environment;
use codex_login::AuthChangeState;
use codex_otel::auth_storage::AuthStorageOriginator;
use codex_protocol::mcp::ClientMcpExtensions;
use codex_protocol::mcp::McpServerInfo;
use codex_protocol::protocol::Event;
@@ -286,6 +287,7 @@ fn codex_apps_reconnect_backoff(consecutive_failures: u32) -> Duration {
#[derive(Clone)]
struct ManagedClientStartup {
originator: AuthStorageOriginator,
server_name: String,
server: EffectiveMcpServer,
store_mode: OAuthCredentialsStoreMode,
@@ -316,6 +318,7 @@ impl ManagedClientStartup {
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = None;
let Self {
originator,
server_name,
server,
store_mode,
@@ -343,7 +346,7 @@ impl ManagedClientStartup {
.startup_timeout_sec
.unwrap_or(DEFAULT_STARTUP_TIMEOUT);
let cancel_token_for_fut = cancel_token;
async move {
let startup = async move {
let tool_catalog_fetch_ticket = tool_catalog_cache_context
.as_ref()
.map(McpToolCatalogCacheContext::begin_fetch);
@@ -420,10 +423,8 @@ impl ManagedClientStartup {
startup_complete.store(true, Ordering::Release);
outcome
}
.in_current_span()
.boxed()
.shared()
};
originator.scope(startup).in_current_span().boxed().shared()
}
}
@@ -480,6 +481,7 @@ impl AsyncManagedClient {
let startup_complete = Arc::new(AtomicBool::new(false));
let server_capabilities = Arc::new(StdMutex::new(None));
let startup = Arc::new(ManagedClientStartup {
originator: AuthStorageOriginator::current(),
server_name,
server,
store_mode,
+12 -6
View File
@@ -11,6 +11,7 @@ use codex_mcp::ElicitationReviewRequest;
use codex_mcp::ElicitationReviewer;
use codex_mcp::ElicitationReviewerHandle;
use codex_mcp::MCP_TOOL_CODEX_APPS_META_KEY;
use codex_otel::auth_storage::AuthStorageOriginator;
use codex_prompts::ResolvedModelMessages;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
@@ -343,7 +344,9 @@ impl Session {
input.mcp_servers.contains_key(CODEX_APPS_MCP_SERVER_NAME),
"unknown MCP server '{CODEX_APPS_MCP_SERVER_NAME}'"
);
let refreshed = self.services.mcp_runtime.replace_fresh(input).await;
let refreshed = AuthStorageOriginator::from_client_name(&desired.originator)
.scope(self.services.mcp_runtime.replace_fresh(input))
.await;
self.services.thread_extension_data.insert(selected_plugins);
refreshed
}
@@ -375,11 +378,14 @@ impl Session {
self.mark_mcp_runtime_dirty();
}
let recovered_oauth_servers = self
.services
.mcp_runtime
.updated_oauth_credentials_after_auth_failure()
.await;
let recovered_oauth_servers =
AuthStorageOriginator::from_client_name(&turn_context.originator)
.scope(
self.services
.mcp_runtime
.updated_oauth_credentials_after_auth_failure(),
)
.await;
if !recovered_oauth_servers.is_empty()
&& let Ok(_refresh) = self.mcp_refresh.acquire().await
&& self
+4 -1
View File
@@ -16,6 +16,7 @@ use codex_mcp::McpServerSource;
use codex_mcp::McpStartupPolicy;
use codex_mcp::PreparedMcpCall;
use codex_mcp::ToolInfo;
use codex_otel::auth_storage::AuthStorageOriginator;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use std::collections::HashSet;
@@ -300,7 +301,9 @@ impl Session {
ready_selected_capability_roots,
elicitation_reviewer,
);
self.services.mcp_runtime.replace(input).await;
AuthStorageOriginator::from_client_name(&desired.originator)
.scope(self.services.mcp_runtime.replace(input))
.await;
self.services.thread_extension_data.insert(selected_plugins);
}
@@ -0,0 +1,177 @@
//! Exercises thread-owned storage attribution through turn-time MCP authentication recovery.
use anyhow::Result;
use codex_config::McpServerConfig;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_core::StartThreadOptions;
use codex_core::TurnInputRequest;
use codex_history::InitialHistory;
use codex_history::RolloutItem;
use codex_otel::MetricsClient;
use codex_otel::MetricsConfig;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::SessionMeta;
use codex_protocol::protocol::SessionMetaLine;
use codex_protocol::user_input::UserInput;
use core_test_support::responses;
use core_test_support::skip_if_no_network;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use opentelemetry_sdk::metrics::data::AggregatedMetrics;
use opentelemetry_sdk::metrics::data::MetricData;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::collections::BTreeMap;
use std::sync::Arc;
use wiremock::Mock;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
fn storage_loads(metrics: &MetricsClient) -> Result<BTreeMap<String, u64>> {
let mut counts = BTreeMap::new();
for metric in metrics
.snapshot()?
.scope_metrics()
.flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics)
{
if metric.name() == "codex.auth_storage.operation"
&& let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data()
{
for point in sum.data_points() {
let tags: BTreeMap<_, _> = point
.attributes()
.map(|tag| (tag.key.to_string(), tag.value.to_string()))
.collect();
if tags.get("credential_kind").map(String::as_str) == Some("mcp")
&& tags.get("operation").map(String::as_str) == Some("load")
{
*counts.entry(tags["originator"].clone()).or_default() += point.value();
}
}
}
}
Ok(counts)
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn recovery_without_a_snapshot_keeps_each_threads_originator() -> Result<()> {
skip_if_no_network!(Ok(()));
// Isolate global metrics and CODEX_HOME from other integration tests.
if std::env::var_os("CODEX_MCP_TELEMETRY_TEST_CHILD").is_none() {
let home = tempfile::tempdir()?;
let output = tokio::process::Command::new(std::env::current_exe()?)
.args(["--exact", "suite::rmcp_client::storage_telemetry_tests::recovery_without_a_snapshot_keeps_each_threads_originator", "--nocapture"])
.env("CODEX_MCP_TELEMETRY_TEST_CHILD", "1")
.env("CODEX_HOME", home.path())
.output().await?;
assert!(
output.status.success(),
"{}{}",
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr)
);
return Ok(());
}
let server = responses::start_mock_server().await;
let mcp = responses::start_mock_server().await;
Mock::given(method("POST"))
.and(path("/mcp"))
.respond_with(ResponseTemplate::new(/*s*/ 401).append_header("www-authenticate", "Bearer"))
.mount(&mcp)
.await;
Mock::given(method("GET"))
.and(path("/.well-known/oauth-authorization-server/mcp"))
.respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({
"authorization_endpoint": format!("{}/authorize", mcp.uri()),
"token_endpoint": format!("{}/token", mcp.uri()),
})))
.mount(&mcp)
.await;
let home = Arc::new(tempfile::tempdir()?);
// The OAuth endpoint remains on the host when the test uses a remote executor.
let fixture = test_codex()
.with_home(home)
.build_with_remote_and_local_env(&server)
.await?;
let server_config: McpServerConfig =
serde_json::from_value(json!({"url": format!("{}/mcp", mcp.uri())}))?;
let mut config = fixture.config.clone();
config.mcp_oauth_credentials_store_mode = OAuthCredentialsStoreMode::File;
config
.mcp_servers
.set([("reauth".into(), server_config)].into())?;
let metrics = codex_otel::install_global_metrics(MetricsClient::new(
MetricsConfig::in_memory("test", "test", "1", Default::default()).with_runtime_reader(),
)?);
let originators = ["codex_desktop", "codex_vscode"];
let mut threads = Vec::new();
for originator in originators {
let thread = fixture
.thread_manager
.start_thread(StartThreadOptions {
initial_history: InitialHistory::Forked(vec![RolloutItem::SessionMeta(
SessionMetaLine {
meta: SessionMeta {
originator: originator.into(),
..Default::default()
},
git: None,
},
)]),
environments: Some(fixture.codex.environment_selections().await),
..StartThreadOptions::new(config.clone())
})
.await?
.thread;
let EventMsg::McpStartupComplete(startup) = wait_for_event(&thread, |event| {
matches!(event, EventMsg::McpStartupComplete(_))
})
.await
else {
unreachable!();
};
assert_eq!(
startup
.failed
.iter()
.map(|failure| failure.server.as_str())
.collect::<Vec<_>>(),
vec!["reauth"]
);
threads.push(thread);
responses::mount_sse_once(
&server,
responses::sse(vec![
responses::ev_assistant_message("msg", "Sign in to continue."),
responses::ev_completed("resp"),
]),
)
.await;
}
// The runtime reader exports deltas; drain startup observations before recovery.
let _ = storage_loads(&metrics)?;
futures::future::try_join_all(threads.iter().map(|thread| async {
thread
.start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
text: "Continue".into(),
text_elements: Vec::new(),
}]))
.await?;
wait_for_event(thread, |event| matches!(event, EventMsg::TurnComplete(_))).await;
Ok::<_, anyhow::Error>(())
}))
.await?;
let mut after = storage_loads(&metrics)?;
after.retain(|_, count| *count > 0);
assert_eq!(
after.keys().map(String::as_str).collect::<Vec<_>>(),
originators
);
for thread in threads {
thread.shutdown_and_wait().await?;
}
fixture.codex.shutdown_and_wait().await?;
metrics.shutdown()?;
Ok(())
}
+2
View File
@@ -132,6 +132,8 @@ use wiremock::MockServer;
#[path = "mcp_oauth_refresh_tests.rs"]
mod oauth_refresh_tests;
#[path = "mcp_storage_telemetry_tests.rs"]
mod storage_telemetry_tests;
static OPENAI_PNG: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAD0AAAA9CAYAAAAeYmHpAAAE6klEQVR4Aeyau44UVxCGx1fZsmRLlm3Zoe0XcGQ5cUiCCIgJeS9CHgAhMkISQnIuGQgJEkBcxLW+nqnZ6uqqc+nuWRC7q/P3qetf9e+MtOwyX25O4Nep6JPyop++0qev9HrfgZ+F6r2DuB/vHOrt/UIkqdDHYvujOW6fO7h/CNEI+a5jc+pBR8uy0jVFsziYu5HtfSUk+Io34q921hLNctFSX0gwww+S8wce8K1LfCU+cYW4888aov8NxqvQILUPPReLOrm6zyLxa4i+6VZuFbJo8d1MOHZm+7VUtB/aIvhPWc/3SWg49JcwFLlHxuXKjtyloo+YNhuW3VS+WPBuUEMvCFKjEDVgFBQHXrnazpqiSxNZCkQ1kYiozsbm9Oz7l4i2Il7vGccGNWAc3XosDrZe/9P3ZnMmzHNEQw4smf8RQ87XEAMsC7Az0Au+dgXerfH4+sHvEc0SYGic8WBBUGqFH2gN7yDrazy7m2pbRTeRmU3+MjZmr1h6LJgPbGy23SI6GlYT0brQ71IY8Us4PNQCm+zepSbaD2BY9xCaAsD9IIj/IzFmKMSdHHonwdZATbTnYREf6/VZGER98N9yCWIvXQwXDoDdhZJoT8jwLnJXDB9w4Sb3e6nK5ndzlkTLnP3JBu4LKkbrYrU69gCVceV0JvpyuW1xlsUVngzhwMetn/XamtTORF9IO5YnWNiyeF9zCAfqR3fUW+vZZKLtgP+ts8BmQRBREAdRDhH3o8QuRh/YucNFz2BEjxbRN6LGzphfKmvP6v6QhqIQyZ8XNJ0W0X83MR1PEcJBNO2KC2Z1TW/v244scp9FwRViZxIOBF0Lctk7ZVSavdLvRlV1hz/ysUi9sr8CIcB3nvWBwA93ykTz18eAYxQ6N/K2DkPA1lv3iXCwmDUT7YkjIby9siXueIJj9H+pzSqJ9oIuJWTUgSSt4WO7o/9GGg0viR4VinNRUDoIj34xoCd6pxD3aK3zfdbnx5v1J3ZNNEJsE0sBG7N27ReDrJc4sFxz7dI/ZAbOmmiKvHBitQXpAdR6+F7v+/ol/tOouUV01EeMZQF2BoQDn6dP4XNr+j9GZEtEK1/L8pFw7bd3a53tsTa7WD+054jOFmPg1XBKPQgnqFfmFcy32ZRvjmiIIQTYFvyDxQ8nH8WIwwGwlyDjDznnilYyFr6njrlZwsKkBpO59A7OwgdzPEWRm+G+oeb7IfyNuzjEEVLrOVxJsxvxwF8kmCM6I2QYmJunz4u4TrADpfl7mlbRTWQ7VmrBzh3+C9f6Grc3YoGN9dg/SXFthpRsT6vobfXRs2VBlgBHXVMLHjDNbIZv1sZ9+X3hB09cXdH1JKViyG0+W9bWZDa/r2f9zAFR71sTzGpMSWz2iI4YssWjWo3REy1MDGjdwe5e0dFSiAC1JakBvu4/CUS8Eh6dqHdU0Or0ioY3W5ClSqDXAy7/6SRfgw8vt4I+tbvvNtFT2kVDhY5+IGb1rCqYaXNF08vSALsXCPmt0kQNqJT1p5eI1mkIV/BxCY1z85lOzeFbPBQHURkkPTlwTYK9gTVE25l84IbFFN+YJDHjdpn0gq6mrHht0dkcjbM4UL9283O5p77GN+SPW/QwVB4IUYg7Or+Kp7naR6qktP98LNF2UxWo9yObPIT9KYg+hK4i56no4rfnM0qeyFf6AwAAAP//trwR3wAAAAZJREFUAwBZ0sR75itw5gAAAABJRU5ErkJggg==";
+8 -3
View File
@@ -69,6 +69,7 @@ webbrowser = { workspace = true }
which = { workspace = true }
[dev-dependencies]
opentelemetry_sdk = { workspace = true, features = ["testing"] }
codex-http-client = { workspace = true }
codex-utils-cargo-bin = { workspace = true }
pretty_assertions = { workspace = true }
@@ -94,10 +95,14 @@ windows-sys = { version = "0.52", features = ["Win32_Storage_FileSystem"] }
[target.'cfg(any(target_os = "freebsd", target_os = "openbsd"))'.dependencies]
keyring = { workspace = true, features = ["sync-secret-service"] }
# This test is compiled through `#[path]` inside the inline `oauth::tests` module. Cargo-shear
# cannot resolve that nested module path and otherwise reports the linked file as unlinked.
# These tests are compiled through `#[path]` inside the inline `oauth::tests` module. Cargo-shear
# cannot resolve those nested module paths and otherwise reports the linked files as unlinked.
[package.metadata.cargo-shear]
ignored-paths = ["src/oauth/tests/persistor_tests.rs", "src/oauth/tests/credential_store_tests.rs"]
ignored-paths = [
"src/oauth/tests/persistor_tests.rs",
"src/oauth/tests/credential_store_tests.rs",
"src/oauth/tests/telemetry_tests.rs",
]
[lib]
doctest = false
@@ -233,7 +233,7 @@ impl KeyringStore for GatedKeyringStore {
async fn refresh_subject_reread_is_cancellable_and_releases_guard_on_failure() -> Result<()> {
let _home = TempCodexHome::new();
let (_server, issuer) = discovery().await;
let store = ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct);
let store = ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct);
let stored = credentials(&issuer, "user", /*expires_at*/ 0);
let snapshot = StoredOAuthCredentialSnapshot::new(stored.clone(), store);
for outcome in [
@@ -312,7 +312,7 @@ async fn refresh_subject_reread_is_cancellable_and_releases_guard_on_failure() -
async fn refresh_subject_rereads_pinned_credentials_after_id_token_expiry() -> Result<()> {
let _home = TempCodexHome::new();
let (_server, issuer) = discovery().await;
let store = ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct);
let store = ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct);
let stored = credentials(&issuer, "user", /*expires_at*/ 0);
let snapshot = StoredOAuthCredentialSnapshot::new(stored.clone(), store);
let keyring = MockKeyringStore::default();
@@ -341,7 +341,7 @@ async fn refresh_subject_rejects_removed_replaced_or_missing_credentials() -> Re
let _home = TempCodexHome::new();
let (server, issuer) = discovery().await;
let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
let store = ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct);
let store = ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct);
let original = credentials(&issuer, "user", /*expires_at*/ 0);
for change in [
"deleted",
@@ -357,7 +357,7 @@ async fn refresh_subject_rejects_removed_replaced_or_missing_credentials() -> Re
let snapshot = StoredOAuthCredentialSnapshot::new(
original.clone(),
if change == "file" {
ResolvedOAuthCredentialStore::File
ResolvedOAuthCredentialStore::file()
} else {
store
},
+91 -57
View File
@@ -25,6 +25,7 @@ mod refresh_transaction;
mod resolved_store;
mod runtime;
mod store_lock;
mod telemetry;
#[cfg(test)]
#[path = "oauth/test_support.rs"]
@@ -35,6 +36,8 @@ use anyhow::Error;
use anyhow::Result;
use codex_config::types::AuthKeyringBackendKind;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_otel::auth_storage::Operation;
use codex_otel::auth_storage::Store;
use codex_secrets::LocalSecretsNamespace;
use codex_secrets::SecretName;
use codex_secrets::SecretScope;
@@ -214,7 +217,7 @@ impl StoredOAuthCredentialSnapshot {
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
) -> Result<Option<StoredOAuthTokens>> {
if self.store == ResolvedOAuthCredentialStore::File
if self.store == ResolvedOAuthCredentialStore::file()
&& store_mode == OAuthCredentialsStoreMode::Auto
{
return Self::for_runtime_refresh(
@@ -227,7 +230,15 @@ impl StoredOAuthCredentialSnapshot {
.map(|snapshot| snapshot.map(|snapshot| snapshot.credentials));
}
let credentials = match self.store.try_load(&DefaultKeyringStore, server_name, url) {
let mut observation = telemetry::resolved(self.store, Operation::Load);
let result = self.store.try_load(&DefaultKeyringStore, server_name, url);
observation.record_load_attempt(telemetry::store(self.store), &result);
if matches!(self.store.backend, resolved_store::Backend::Keyring(_))
&& let Err(error) = &result
{
telemetry::record_secure_error(&mut observation, error.as_ref());
}
let credentials = match result {
Ok(credentials) => credentials,
Err(error) if oauth_store_is_contended(&error) => return Ok(None),
Err(error) => return Err(error),
@@ -393,7 +404,7 @@ fn load_oauth_tokens_from_direct_keyring<K: KeyringStore>(
Ok(Some(tokens))
}
Ok(None) => Ok(None),
Err(error) => Err(Error::new(error.into_error())),
Err(error) => Err(Error::new(std::io::Error::from(error))),
}
}
@@ -440,7 +451,7 @@ enum OAuthKeyringLoadError {
#[error(transparent)]
StoreLock(#[from] OAuthStoreLockFailure),
/// The selected keyring backend itself was unavailable or its data was invalid.
#[error(transparent)]
#[error("{0}")]
Backend(#[from] anyhow::Error),
}
@@ -463,6 +474,7 @@ pub(crate) fn save_oauth_tokens_with_lock_held(
keyring_backend_kind: AuthKeyringBackendKind,
) -> Result<()> {
let keyring_store = DefaultKeyringStore;
let mut observation = telemetry::policy(store_mode, keyring_backend_kind, Operation::Save);
match store_mode {
OAuthCredentialsStoreMode::Auto => save_oauth_tokens_with_keyring_with_fallback_to_file(
&keyring_store,
@@ -470,13 +482,25 @@ pub(crate) fn save_oauth_tokens_with_lock_held(
server_name,
tokens,
),
OAuthCredentialsStoreMode::File => save_oauth_tokens_to_file(tokens),
OAuthCredentialsStoreMode::Keyring => save_oauth_tokens_with_keyring_and_cleanup_file(
&keyring_store,
keyring_backend_kind,
server_name,
tokens,
),
OAuthCredentialsStoreMode::File => {
let result = save_oauth_tokens_to_file(tokens);
observation.record_save_attempt(Store::File, &result);
result
}
OAuthCredentialsStoreMode::Keyring => {
let result = save_oauth_tokens_with_keyring_and_cleanup_file(
&keyring_store,
store_mode,
keyring_backend_kind,
server_name,
tokens,
);
observation.record_save_attempt(telemetry::keyring(keyring_backend_kind), &result);
if let Err(error) = &result {
telemetry::record_secure_error(&mut observation, error.as_ref());
}
result
}
}
}
@@ -506,17 +530,10 @@ fn save_oauth_tokens_to_direct_keyring<K: KeyringStore>(
let serialized = serde_json::to_string(tokens).context("failed to serialize OAuth tokens")?;
let key = compute_store_key(server_name, &tokens.url)?;
match keyring_store.save(KEYRING_SERVICE, &key, &serialized) {
Ok(()) => Ok(()),
Err(error) => {
let message = format!(
"failed to write OAuth tokens to keyring: {}",
error.message()
);
warn!("{message}");
Err(Error::new(error.into_error()).context(message))
}
}
keyring_store
.save(KEYRING_SERVICE, &key, &serialized)
.map_err(std::io::Error::from)
.context("failed to write OAuth tokens to keyring")
}
/// Saves one credential while holding the Secrets aggregate-store lock across the mutation.
@@ -558,20 +575,20 @@ fn save_oauth_tokens_to_secrets_keyring_with_lock_held<K: KeyringStore + Clone +
/// Saves to the selected keyring backend, then best-effort removes the fallback File entry.
fn save_oauth_tokens_with_keyring_and_cleanup_file<K: KeyringStore + Clone + 'static>(
keyring_store: &K,
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
server_name: &str,
tokens: &StoredOAuthTokens,
) -> Result<()> {
save_oauth_tokens_with_keyring(keyring_store, keyring_backend_kind, server_name, tokens)?;
let key = compute_store_key(server_name, &tokens.url)?;
if let Err(error) = delete_oauth_tokens_from_file(&key) {
warn!(
server_name,
keyring_backend = ?keyring_backend_kind,
error = %error,
"failed to remove OAuth tokens from fallback storage"
);
let mut observation = telemetry::policy(store_mode, keyring_backend_kind, Operation::Cleanup);
let result = delete_oauth_tokens_from_file(&key);
observation.record_delete_attempt(Store::File, &result);
if result.is_err() {
warn!("failed to remove OAuth tokens from fallback storage");
}
drop(observation);
Ok(())
}
@@ -581,22 +598,34 @@ fn save_oauth_tokens_with_keyring_with_fallback_to_file<K: KeyringStore + Clone
server_name: &str,
tokens: &StoredOAuthTokens,
) -> Result<()> {
match save_oauth_tokens_with_keyring_and_cleanup_file(
let mut observation = telemetry::policy(
OAuthCredentialsStoreMode::Auto,
keyring_backend_kind,
Operation::Save,
);
let result = save_oauth_tokens_with_keyring_and_cleanup_file(
keyring_store,
OAuthCredentialsStoreMode::Auto,
keyring_backend_kind,
server_name,
tokens,
) {
);
observation.record_save_attempt(telemetry::keyring(keyring_backend_kind), &result);
match result {
Ok(()) => Ok(()),
// As on load, a store lock failure is a coordination failure rather than evidence that
// the keyring backend is unavailable. Falling back could leave a newer File token hidden
// behind a stale Secrets entry.
Err(error) if error.downcast_ref::<OAuthStoreLockFailure>().is_some() => Err(error),
Err(error) if error.downcast_ref::<OAuthStoreLockFailure>().is_some() => {
telemetry::record_secure_error(&mut observation, error.as_ref());
Err(error)
}
Err(error) => {
let message = error.to_string();
warn!("falling back to file storage for OAuth tokens: {message}");
save_oauth_tokens_to_file(tokens)
.with_context(|| format!("failed to write OAuth tokens to keyring: {message}"))
telemetry::record_secure_error(&mut observation, error.as_ref());
let result = save_oauth_tokens_to_file(tokens);
observation.record_save_attempt(Store::File, &result);
result.with_context(|| format!("failed to write OAuth tokens to keyring: {message}"))
}
}
}
@@ -636,25 +665,28 @@ fn delete_oauth_tokens_from_keyring_and_file<K: KeyringStore + Clone + 'static>(
server_name: &str,
url: &str,
) -> Result<bool> {
let key = compute_store_key(server_name, url)?;
let keyring_result =
delete_oauth_tokens_from_keyring(keyring_store, keyring_backend_kind, server_name, url);
let keyring_removed = match keyring_result {
Ok(removed) => removed,
Err(error) => {
let message = error.to_string();
warn!("failed to delete OAuth tokens from keyring: {message}");
match store_mode {
OAuthCredentialsStoreMode::Auto | OAuthCredentialsStoreMode::Keyring => {
return Err(error).context("failed to delete OAuth tokens from keyring");
let mut observation = telemetry::policy(store_mode, keyring_backend_kind, Operation::Delete);
let result = compute_store_key(server_name, url).and_then(|key| {
let keyring_result =
delete_oauth_tokens_from_keyring(keyring_store, keyring_backend_kind, server_name, url);
let keyring_removed = match keyring_result {
Ok(removed) => removed,
Err(error) => {
warn!("failed to delete OAuth tokens from keyring");
match store_mode {
OAuthCredentialsStoreMode::Auto | OAuthCredentialsStoreMode::Keyring => {
return Err(error).context("failed to delete OAuth tokens from keyring");
}
OAuthCredentialsStoreMode::File => false,
}
OAuthCredentialsStoreMode::File => false,
}
}
};
};
let file_removed = delete_oauth_tokens_from_file(&key)?;
Ok(keyring_removed || file_removed)
let file_removed = delete_oauth_tokens_from_file(&key)?;
Ok(keyring_removed || file_removed)
});
observation.record_delete_attempt(Store::Multiple, &result);
result
}
fn delete_oauth_tokens_from_keyring<K: KeyringStore + Clone + 'static>(
@@ -685,7 +717,7 @@ fn delete_oauth_tokens_from_direct_keyring<K: KeyringStore>(
let key = compute_store_key(server_name, url)?;
keyring_store
.delete(KEYRING_SERVICE, &key)
.map_err(|error| Error::new(error.into_error()))
.map_err(|error| Error::new(std::io::Error::from(error)))
}
fn delete_oauth_tokens_from_secrets_keyring<K: KeyringStore + Clone + 'static>(
@@ -786,7 +818,7 @@ impl OAuthPersistor {
expires_at,
};
if last_credentials.as_ref() != Some(&stored) {
self.inner.credential_store.save(
self.inner.credential_store.save_with_refresh_telemetry(
&DefaultKeyringStore,
&self.inner.server_name,
&stored,
@@ -1160,6 +1192,8 @@ mod tests {
mod credential_store_tests;
#[path = "persistor_tests.rs"]
mod persistor_tests;
#[path = "telemetry_tests.rs"]
mod telemetry_tests;
use super::test_support::TempCodexHome;
@@ -1224,7 +1258,7 @@ mod tests {
.expect("tokens should load from keyring");
assert_eq!(
resolved.store,
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct)
ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct)
);
assert_tokens_match_without_expiry(&resolved.tokens, &expected);
Ok(())
@@ -1247,7 +1281,7 @@ mod tests {
AuthKeyringBackendKind::Direct,
)?
.expect("tokens should load from fallback");
assert_eq!(resolved.store, ResolvedOAuthCredentialStore::File);
assert_eq!(resolved.store, ResolvedOAuthCredentialStore::file());
assert_tokens_match_without_expiry(&resolved.tokens, &expected);
Ok(())
}
@@ -1271,7 +1305,7 @@ mod tests {
AuthKeyringBackendKind::Direct,
)?
.expect("tokens should load from fallback");
assert_eq!(resolved.store, ResolvedOAuthCredentialStore::File);
assert_eq!(resolved.store, ResolvedOAuthCredentialStore::file());
assert_tokens_match_without_expiry(&resolved.tokens, &expected);
Ok(())
}
@@ -1298,7 +1332,7 @@ mod tests {
)?;
assert_eq!(fs::read(fallback_path)?, fallback_before);
let loaded = ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct)
let loaded = ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct)
.load(&store, &keyring_tokens.server_name, &keyring_tokens.url)?
.expect("tokens should load from the selected keyring store");
assert_tokens_match_without_expiry(&loaded, &keyring_tokens);
@@ -243,9 +243,11 @@ impl<K: KeyringStore + Clone + 'static> CredentialStore for OAuthCredentialStore
.map_err(credential_store_error)?;
let tokens = tokio::task::spawn_blocking(move || -> Result<StoredOAuthTokens> {
let _guard = guard;
inner
.store
.save(&inner.keyring, &inner.server_name, &tokens)?;
inner.store.save_with_refresh_telemetry(
&inner.keyring,
&inner.server_name,
&tokens,
)?;
Ok(tokens)
})
.await
@@ -45,7 +45,7 @@ impl StoredOAuthCredentialSnapshot {
&self,
keyring_store: &K,
) -> Result<StoredOAuthTokens> {
if !matches!(self.store, ResolvedOAuthCredentialStore::Keyring(_)) {
if self.store == ResolvedOAuthCredentialStore::file() {
bail!("enterprise IdP credentials require keyring storage");
}
let previous = &self.credentials;
@@ -25,7 +25,7 @@ async fn keyring_failure_does_not_reuse_the_pinned_refresh_token() -> Result<()>
let _lock = RefreshCredentialLock::acquire_for_server(&tokens.server_name, &tokens.url).await?;
let snapshot = StoredOAuthCredentialSnapshot::new(
tokens,
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct),
);
let keyring = MockKeyringStore::default();
keyring.set_error(
@@ -64,7 +64,7 @@ fn ordinary_oauth_names_cannot_alias_enterprise_credential_keys() -> Result<()>
);
let keyring = MockKeyringStore::default();
let store = ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct);
let store = ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct);
let enterprise_tokens: StoredOAuthTokens = serde_json::from_value(json!({
"server_name": enterprise_name, "url": issuer, "issuer": issuer,
"client_id": "idp-client", "token_response": {
@@ -246,13 +246,13 @@ impl OAuthPersistor {
// TODO: Add a bounded persistence retry only if telemetry shows this is common; never
// silently switch stores or continue with an unpersisted credential.
debug!("persisting refreshed MCP OAuth credentials to the resolved store");
if let Err(error) =
self.inner
.credential_store
.save(keyring_store, &self.inner.server_name, &refreshed)
{
let result = self.inner.credential_store.save_with_refresh_telemetry(
keyring_store,
&self.inner.server_name,
&refreshed,
);
if let Err(error) = result {
warn!(
error = %error,
"failed to persist refreshed MCP OAuth credentials; returning the error and restoring the previous in-process credentials"
);
install_tokens_in_manager(&mut guard, &latest, self.inner.oauth_config.as_ref())
+186 -64
View File
@@ -1,11 +1,15 @@
//! Resolves the configured MCP OAuth store and pins that concrete source for one client lifecycle.
//! Retains the originating client for later background storage observations.
use super::telemetry;
use anyhow::Context;
use anyhow::Result;
use codex_config::types::AuthKeyringBackendKind;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_keyring_store::KeyringStore;
use tracing::warn;
use codex_otel::auth_storage::AuthStorageOriginator;
use codex_otel::auth_storage::Operation;
use codex_otel::auth_storage::Store;
use super::OAuthKeyringLoadError;
use super::OAuthStore;
@@ -29,13 +33,57 @@ use super::save_oauth_tokens_with_keyring;
/// client that loaded credentials from one store must reread, refresh, persist, and remove only
/// through that store. A mid-lifecycle backend failure is unexpected and must return an error
/// rather than falling back to another possibly stale refresh token.
#[derive(Debug, Clone, Copy)]
pub(crate) struct ResolvedOAuthCredentialStore {
pub(super) backend: Backend,
pub(super) mode: OAuthCredentialsStoreMode,
pub(super) kind: AuthKeyringBackendKind,
pub(super) originator: AuthStorageOriginator,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ResolvedOAuthCredentialStore {
pub(super) enum Backend {
File,
Keyring(AuthKeyringBackendKind),
}
// Telemetry metadata must not change credential identity or trigger a reconnect.
impl PartialEq for ResolvedOAuthCredentialStore {
fn eq(&self, other: &Self) -> bool {
self.backend == other.backend
}
}
impl Eq for ResolvedOAuthCredentialStore {}
impl ResolvedOAuthCredentialStore {
pub(crate) fn file() -> Self {
Self {
backend: Backend::File,
mode: OAuthCredentialsStoreMode::File,
kind: AuthKeyringBackendKind::Direct,
originator: AuthStorageOriginator::current(),
}
}
pub(crate) fn keyring(kind: AuthKeyringBackendKind) -> Self {
Self {
backend: Backend::Keyring(kind),
mode: OAuthCredentialsStoreMode::Keyring,
kind,
originator: AuthStorageOriginator::current(),
}
}
fn with_policy(
mut self,
mode: OAuthCredentialsStoreMode,
kind: AuthKeyringBackendKind,
) -> Self {
self.mode = mode;
self.kind = kind;
self
}
/// Loads credentials only from this already-resolved authority.
///
/// Unlike `resolve_oauth_tokens_from_store_policy`, this never evaluates configured
@@ -46,10 +94,11 @@ impl ResolvedOAuthCredentialStore {
server_name: &str,
url: &str,
) -> Result<Option<StoredOAuthTokens>> {
match self {
Self::File => load_oauth_tokens_from_file(server_name, url)
let mut observation = telemetry::resolved(self, Operation::Load);
let result = match self.backend {
Backend::File => load_oauth_tokens_from_file(server_name, url)
.context("failed to reread OAuth tokens from resolved file storage"),
Self::Keyring(keyring_backend_kind) => load_oauth_tokens_from_keyring(
Backend::Keyring(keyring_backend_kind) => load_oauth_tokens_from_keyring(
keyring_store,
keyring_backend_kind,
server_name,
@@ -59,7 +108,14 @@ impl ResolvedOAuthCredentialStore {
.context(
"failed to reread OAuth tokens from resolved keyring storage; refusing file fallback",
),
};
observation.record_load_attempt(telemetry::store(self), &result);
if matches!(self.backend, Backend::Keyring(_))
&& let Err(error) = &result
{
telemetry::record_secure_error(&mut observation, error.as_ref());
}
result
}
/// Reads the selected authority without waiting for its aggregate-store lock.
@@ -69,16 +125,18 @@ impl ResolvedOAuthCredentialStore {
server_name: &str,
url: &str,
) -> Result<Option<StoredOAuthTokens>> {
match self {
Self::File => {
match self.backend {
Backend::File => {
let _store_lock = OAuthStoreLock::try_acquire_for_read(OAuthStore::File)?;
load_oauth_tokens_from_file_with_lock_held(server_name, url)
.context("failed to probe OAuth tokens from resolved file storage")
}
Self::Keyring(AuthKeyringBackendKind::Direct) => {
self.load(keyring_store, server_name, url)
Backend::Keyring(AuthKeyringBackendKind::Direct) => {
load_oauth_tokens_from_keyring(keyring_store, AuthKeyringBackendKind::Direct, server_name, url)
.map_err(anyhow::Error::from)
.context("failed to reread OAuth tokens from resolved keyring storage; refusing file fallback")
}
Self::Keyring(AuthKeyringBackendKind::Secrets) => {
Backend::Keyring(AuthKeyringBackendKind::Secrets) => {
let _store_lock = OAuthStoreLock::try_acquire_for_read(OAuthStore::Secrets)?;
load_oauth_tokens_from_secrets_keyring_with_lock_held(
keyring_store,
@@ -97,15 +155,41 @@ impl ResolvedOAuthCredentialStore {
server_name: &str,
tokens: &StoredOAuthTokens,
) -> Result<()> {
match self {
Self::File => save_oauth_tokens_to_file(tokens),
Self::Keyring(keyring_backend_kind) => save_oauth_tokens_with_keyring(
let mut observation = telemetry::resolved(self, Operation::Save);
let result = match self.backend {
Backend::File => save_oauth_tokens_to_file(tokens),
Backend::Keyring(keyring_backend_kind) => save_oauth_tokens_with_keyring(
keyring_store,
keyring_backend_kind,
server_name,
tokens,
),
};
observation.record_save_attempt(telemetry::store(self), &result);
if matches!(self.backend, Backend::Keyring(_))
&& let Err(error) = &result
{
telemetry::record_secure_error(&mut observation, error.as_ref());
}
result
}
/// Records the refresh-persistence milestone in addition to the underlying pinned save.
pub(crate) fn save_with_refresh_telemetry<K: KeyringStore + Clone + 'static>(
self,
keyring_store: &K,
server_name: &str,
tokens: &StoredOAuthTokens,
) -> Result<()> {
let mut observation = telemetry::resolved(self, Operation::RefreshPersist);
let result = self.save(keyring_store, server_name, tokens);
observation.record_save_attempt(telemetry::store(self), &result);
if matches!(self.backend, Backend::Keyring(_))
&& let Err(error) = &result
{
telemetry::record_secure_error(&mut observation, error.as_ref());
}
result
}
/// Deletes credentials only from this already-resolved authority.
@@ -115,18 +199,24 @@ impl ResolvedOAuthCredentialStore {
server_name: &str,
url: &str,
) -> Result<bool> {
match self {
Self::File => {
let key = compute_store_key(server_name, url)?;
delete_oauth_tokens_from_file(&key)
}
Self::Keyring(AuthKeyringBackendKind::Direct) => {
let mut observation = telemetry::resolved(self, Operation::Delete);
let result = match self.backend {
Backend::File => compute_store_key(server_name, url)
.and_then(|key| delete_oauth_tokens_from_file(&key)),
Backend::Keyring(AuthKeyringBackendKind::Direct) => {
delete_oauth_tokens_from_direct_keyring(keyring_store, server_name, url)
}
Self::Keyring(AuthKeyringBackendKind::Secrets) => {
Backend::Keyring(AuthKeyringBackendKind::Secrets) => {
delete_oauth_tokens_from_secrets_keyring(keyring_store, server_name, url)
}
};
observation.record_delete_attempt(telemetry::store(self), &result);
if matches!(self.backend, Backend::Keyring(_))
&& let Err(error) = &result
{
telemetry::record_secure_error(&mut observation, error.as_ref());
}
result
}
}
@@ -143,6 +233,7 @@ pub(crate) fn resolve_oauth_tokens_from_store_policy<K: KeyringStore + Clone + '
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
) -> Result<Option<ResolvedOAuthTokens>> {
let mut observation = telemetry::policy(store_mode, keyring_backend_kind, Operation::Load);
match store_mode {
OAuthCredentialsStoreMode::Auto => {
// Auto remains keyring-first at lifecycle startup. The returned source is then pinned
@@ -152,59 +243,80 @@ pub(crate) fn resolve_oauth_tokens_from_store_policy<K: KeyringStore + Clone + '
// when keyring availability differs. Solving that safely requires durable backend
// selection or reconciliation of legacy entries and is intentionally outside this
// stack.
match load_oauth_tokens_from_keyring(
let result = load_oauth_tokens_from_keyring(
keyring_store,
keyring_backend_kind,
server_name,
url,
) {
);
observation.record_load_attempt(telemetry::keyring(keyring_backend_kind), &result);
match result {
Ok(Some(tokens)) => Ok(Some(ResolvedOAuthTokens {
tokens,
store: ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind),
store: ResolvedOAuthCredentialStore::keyring(keyring_backend_kind)
.with_policy(store_mode, keyring_backend_kind),
})),
Ok(None) => Ok(
load_oauth_tokens_from_file(server_name, url)?.map(|tokens| {
ResolvedOAuthTokens {
tokens,
store: ResolvedOAuthCredentialStore::File,
}
}),
),
Ok(None) => {
let result = load_oauth_tokens_from_file(server_name, url);
observation.record_load_attempt(Store::File, &result);
Ok(result?.map(|tokens| ResolvedOAuthTokens {
tokens,
store: ResolvedOAuthCredentialStore::file()
.with_policy(store_mode, keyring_backend_kind),
}))
}
// Auto may fall back when the keyring backend is unavailable, but a Secrets
// aggregate-lock failure means authority may be changing. Consulting File in
// that state could replay credentials hidden behind a newer Secrets entry.
Err(OAuthKeyringLoadError::StoreLock(error)) => Err(error.into()),
Err(OAuthKeyringLoadError::StoreLock(error)) => {
telemetry::record_secure_error(&mut observation, &error);
Err(error.into())
}
Err(error) => {
warn!("failed to read OAuth tokens from keyring: {error}");
Ok(load_oauth_tokens_from_file(server_name, url)
telemetry::record_secure_error(&mut observation, &error);
let result = load_oauth_tokens_from_file(server_name, url);
observation.record_load_attempt(Store::File, &result);
Ok(result
.with_context(|| {
format!("failed to read OAuth tokens from keyring: {error}")
})?
.map(|tokens| ResolvedOAuthTokens {
tokens,
store: ResolvedOAuthCredentialStore::File,
store: ResolvedOAuthCredentialStore::file()
.with_policy(store_mode, keyring_backend_kind),
}))
}
}
}
OAuthCredentialsStoreMode::File => Ok(load_oauth_tokens_from_file(server_name, url)?.map(
|tokens| ResolvedOAuthTokens {
OAuthCredentialsStoreMode::File => {
let result = load_oauth_tokens_from_file(server_name, url);
observation.record_load_attempt(Store::File, &result);
Ok(result?.map(|tokens| ResolvedOAuthTokens {
tokens,
store: ResolvedOAuthCredentialStore::File,
},
)),
OAuthCredentialsStoreMode::Keyring => Ok(load_oauth_tokens_from_keyring(
keyring_store,
keyring_backend_kind,
server_name,
url,
)
.map_err(anyhow::Error::from)
.context("failed to read OAuth tokens from keyring")?
.map(|tokens| ResolvedOAuthTokens {
tokens,
store: ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind),
})),
store: ResolvedOAuthCredentialStore::file()
.with_policy(store_mode, keyring_backend_kind),
}))
}
OAuthCredentialsStoreMode::Keyring => {
let result = load_oauth_tokens_from_keyring(
keyring_store,
keyring_backend_kind,
server_name,
url,
);
observation.record_load_attempt(telemetry::keyring(keyring_backend_kind), &result);
if let Err(error) = &result {
telemetry::record_secure_error(&mut observation, error);
}
Ok(result
.map_err(anyhow::Error::from)
.context("failed to read OAuth tokens from keyring")?
.map(|tokens| ResolvedOAuthTokens {
tokens,
store: ResolvedOAuthCredentialStore::keyring(keyring_backend_kind)
.with_policy(store_mode, keyring_backend_kind),
}))
}
}
}
@@ -215,24 +327,34 @@ pub(crate) fn try_resolve_oauth_tokens_from_store_policy<K: KeyringStore + Clone
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
) -> Result<Option<ResolvedOAuthTokens>> {
let load = |store: ResolvedOAuthCredentialStore| {
store
.try_load(keyring_store, server_name, url)
.map(|tokens| tokens.map(|tokens| ResolvedOAuthTokens { tokens, store }))
let mut observation = telemetry::policy(store_mode, keyring_backend_kind, Operation::Load);
let mut load = |store: ResolvedOAuthCredentialStore| {
let result = store.try_load(keyring_store, server_name, url);
observation.record_load_attempt(telemetry::store(store), &result);
if matches!(store.backend, Backend::Keyring(_))
&& let Err(error) = &result
{
telemetry::record_secure_error(&mut observation, error.as_ref());
}
result.map(|tokens| tokens.map(|tokens| ResolvedOAuthTokens { tokens, store }))
};
let keyring = ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind);
let keyring = ResolvedOAuthCredentialStore::keyring(keyring_backend_kind)
.with_policy(store_mode, keyring_backend_kind);
match store_mode {
OAuthCredentialsStoreMode::File => load(ResolvedOAuthCredentialStore::File),
OAuthCredentialsStoreMode::File => {
load(ResolvedOAuthCredentialStore::file().with_policy(store_mode, keyring_backend_kind))
}
OAuthCredentialsStoreMode::Keyring => load(keyring),
OAuthCredentialsStoreMode::Auto => match load(keyring) {
Ok(Some(tokens)) => Ok(Some(tokens)),
Ok(None) => load(ResolvedOAuthCredentialStore::File),
Ok(None) => load(
ResolvedOAuthCredentialStore::file().with_policy(store_mode, keyring_backend_kind),
),
Err(error) if error.downcast_ref::<OAuthStoreLockFailure>().is_some() => Err(error),
Err(error) => {
warn!("failed to read OAuth tokens from keyring: {error}");
load(ResolvedOAuthCredentialStore::File)
.with_context(|| format!("failed to read OAuth tokens from keyring: {error}"))
}
Err(error) => load(
ResolvedOAuthCredentialStore::file().with_policy(store_mode, keyring_backend_kind),
)
.with_context(|| format!("failed to read OAuth tokens from keyring: {error}")),
},
}
}
@@ -0,0 +1,97 @@
//! Maps MCP storage policy, pinned stores, and typed failures to credential-free metric dimensions.
use super::OAuthKeyringLoadError;
use super::OAuthStoreLockFailure;
use super::ResolvedOAuthCredentialStore;
use super::resolved_store::Backend;
use codex_config::types::AuthKeyringBackendKind;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_otel::auth_storage::AuthStorageOriginator;
use codex_otel::auth_storage::CredentialKind;
use codex_otel::auth_storage::Operation;
use codex_otel::auth_storage::StoragePhase;
use codex_otel::auth_storage::StorageTelemetry;
use codex_otel::auth_storage::Store;
use codex_otel::auth_storage::StoreMode;
use std::error::Error;
use std::io::ErrorKind;
/// Preserve aggregate-lock categories even when a transparent load error hides their source.
pub(super) fn record_secure_error(
observation: &mut StorageTelemetry,
error: &(dyn Error + 'static),
) {
let mut cause = Some(error);
while let Some(current) = cause {
let lock = match current.downcast_ref::<OAuthKeyringLoadError>() {
Some(OAuthKeyringLoadError::StoreLock(lock)) => Some(lock),
Some(OAuthKeyringLoadError::Backend(_)) | None => {
current.downcast_ref::<OAuthStoreLockFailure>()
}
};
if let Some(OAuthStoreLockFailure::Timeout {
acquire_timeout, ..
}) = lock
{
let kind = if acquire_timeout.is_zero() {
ErrorKind::WouldBlock
} else {
ErrorKind::TimedOut
};
observation.record_secure_error(&std::io::Error::from(kind));
return;
}
cause = current.source();
}
observation.record_secure_error(error);
}
pub(super) fn keyring(kind: AuthKeyringBackendKind) -> Store {
match kind {
AuthKeyringBackendKind::Direct => Store::DirectKeyring,
AuthKeyringBackendKind::Secrets => Store::Secrets,
}
}
pub(super) fn store(store: ResolvedOAuthCredentialStore) -> Store {
match store.backend {
Backend::File => Store::File,
Backend::Keyring(kind) => keyring(kind),
}
}
pub(super) fn policy(
mode: OAuthCredentialsStoreMode,
kind: AuthKeyringBackendKind,
operation: Operation,
) -> StorageTelemetry {
StorageTelemetry::new(
CredentialKind::Mcp,
store_mode(mode),
keyring(kind),
operation,
AuthStorageOriginator::current(),
)
}
pub(super) fn resolved(
source: ResolvedOAuthCredentialStore,
operation: Operation,
) -> StorageTelemetry {
StorageTelemetry::new(
CredentialKind::Mcp,
store_mode(source.mode),
keyring(source.kind),
operation,
source.originator,
)
.with_phase(StoragePhase::Pinned)
}
fn store_mode(mode: OAuthCredentialsStoreMode) -> StoreMode {
match mode {
OAuthCredentialsStoreMode::File => StoreMode::File,
OAuthCredentialsStoreMode::Auto => StoreMode::Auto,
OAuthCredentialsStoreMode::Keyring => StoreMode::Keyring,
}
}
@@ -36,7 +36,7 @@ async fn mutations_require_and_retain_the_transaction_guard() -> Result<()> {
save_oauth_tokens_to_file(&initial)?;
let store = OAuthCredentialStore::new(
initial.clone(),
ResolvedOAuthCredentialStore::File,
ResolvedOAuthCredentialStore::file(),
MockKeyringStore::default(),
/*oauth_config*/ None,
);
@@ -97,7 +97,7 @@ async fn save_publishes_only_persisted_credentials() -> Result<()> {
let _env = TempCodexHome::new();
let initial = sample_tokens();
let keyring = MockKeyringStore::default();
let authority = ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct);
let authority = ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct);
authority.save(&keyring, &initial.server_name, &initial)?;
let mut fallback = initial.clone();
fallback
@@ -167,7 +167,7 @@ async fn pinned_read_failure_does_not_adopt_fallback_credentials() -> Result<()>
keyring.set_error(&key, KeyringError::Invalid("test".into(), "load".into()));
let store = OAuthCredentialStore::new(
initial.clone(),
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct),
keyring,
/*oauth_config*/ None,
);
@@ -194,7 +194,7 @@ async fn replacement_or_removal_does_not_acknowledge_a_new_runtime_snapshot() ->
save_oauth_tokens_to_file(&initial)?;
let store = OAuthCredentialStore::new(
initial.clone(),
ResolvedOAuthCredentialStore::File,
ResolvedOAuthCredentialStore::file(),
MockKeyringStore::default(),
/*oauth_config*/ None,
);
@@ -238,7 +238,7 @@ async fn storage_roundtrip_preserves_absolute_and_unknown_expiry() -> Result<()>
save_oauth_tokens_to_file(&initial)?;
let store = OAuthCredentialStore::new(
initial.clone(),
ResolvedOAuthCredentialStore::File,
ResolvedOAuthCredentialStore::file(),
MockKeyringStore::default(),
/*oauth_config*/ None,
);
@@ -179,7 +179,7 @@ async fn coordinated_401_refresh_rereads_and_persists_before_retry() -> Result<(
let mut manager = authorization_manager_for(&initial).await?;
manager.set_credential_store(OAuthCredentialStore::new(
initial.clone(),
ResolvedOAuthCredentialStore::File,
ResolvedOAuthCredentialStore::file(),
DefaultKeyringStore,
/*oauth_config*/ None,
));
@@ -396,7 +396,7 @@ async fn resolved_keyring_read_error_preserves_in_memory_credentials() -> Result
initial.server_name.clone(),
initial.url.clone(),
Arc::clone(&manager),
ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct),
Some(initial.clone()),
/*oauth_config*/ None,
);
@@ -875,7 +875,7 @@ async fn coordinated_manager_for(
let mut manager = authorization_manager_for(tokens).await?;
let store = OAuthCredentialStore::new(
tokens.clone(),
ResolvedOAuthCredentialStore::File,
ResolvedOAuthCredentialStore::file(),
DefaultKeyringStore,
/*oauth_config*/ None,
);
@@ -888,13 +888,13 @@ async fn persistor_for(tokens: &StoredOAuthTokens) -> Result<OAuthPersistor> {
tokens.server_name.clone(),
tokens.url.clone(),
Arc::new(TokioMutex::new(authorization_manager_for(tokens).await?)),
ResolvedOAuthCredentialStore::File,
ResolvedOAuthCredentialStore::file(),
Some(tokens.clone()),
/*oauth_config*/ None,
))
}
async fn test_context() -> Result<(TempCodexHome, MockServer, StoredOAuthTokens)> {
pub(super) async fn test_context() -> Result<(TempCodexHome, MockServer, StoredOAuthTokens)> {
let env = TempCodexHome::new();
let server = MockServer::start().await;
mount_oauth_metadata(&server).await;
@@ -902,7 +902,9 @@ async fn test_context() -> Result<(TempCodexHome, MockServer, StoredOAuthTokens)
Ok((env, server, tokens))
}
async fn authorization_manager_for(tokens: &StoredOAuthTokens) -> Result<AuthorizationManager> {
pub(super) async fn authorization_manager_for(
tokens: &StoredOAuthTokens,
) -> Result<AuthorizationManager> {
let oauth_http_client = Arc::new(OAuthHttpClientAdapter::new(
Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new(
OutboundProxyPolicy::ReqwestDefault,
@@ -176,7 +176,7 @@ fn legacy_rmcp_oauth_keyring_credentials_remain_readable() -> Result<()> {
assert_eq!(
resolved.store,
crate::oauth::ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct)
crate::oauth::ResolvedOAuthCredentialStore::keyring(AuthKeyringBackendKind::Direct)
);
assert_tokens_match_without_expiry(&resolved.tokens, &expected);
assert!(crate::oauth::oauth_tokens_are_usable(&resolved.tokens));
@@ -0,0 +1,460 @@
//! Exercises emitted metrics through MCP storage policy and pinned-store operations.
use super::persistor_tests::authorization_manager_for;
use super::persistor_tests::test_context;
use super::*;
use codex_otel::auth_storage::AuthStorageOriginator;
use opentelemetry_sdk::metrics::data::AggregatedMetrics;
use opentelemetry_sdk::metrics::data::MetricData;
use pretty_assertions::assert_eq;
use std::collections::BTreeMap;
type MetricRecord = (String, u64, BTreeMap<String, String>);
#[test]
fn policy_and_pinned_operations_keep_distinct_outcomes_and_originator() -> Result<()> {
if isolated_process()? {
return Ok(());
}
let _env = TempCodexHome::new();
keyring::set_default_credential_builder(keyring::mock::default_credential_builder());
let metrics = test_metrics()?;
let store = MockKeyringStore::default();
let tokens = sample_tokens();
let key = compute_store_key(&tokens.server_name, &tokens.url)?;
store.set_error(
&key,
KeyringError::PlatformFailure(std::io::Error::from(std::io::ErrorKind::Unsupported).into()),
);
save_oauth_tokens_with_keyring_with_fallback_to_file(
&store,
AuthKeyringBackendKind::Direct,
&tokens.server_name,
&tokens,
)?;
store.save(KEYRING_SERVICE, &key, &serde_json::to_string(&tokens)?)?;
let resolved = AuthStorageOriginator::from_client_name("codex_vscode")
.sync_scope(|| {
resolve_oauth_tokens_from_store_policy(
&store,
&tokens.server_name,
&tokens.url,
OAuthCredentialsStoreMode::Auto,
AuthKeyringBackendKind::Direct,
)
})?
.unwrap();
let snapshot = StoredOAuthCredentialSnapshot::new(resolved.tokens, resolved.store);
store.set_error(
&key,
KeyringError::PlatformFailure(std::io::Error::from(std::io::ErrorKind::TimedOut).into()),
);
// Pinned operations retain the original Auto policy and client on another thread.
std::thread::scope(|scope| {
scope
.spawn(|| {
// The SDK mock is empty; this probe must emit exactly one pinned observation.
assert_eq!(
snapshot.reload(
&tokens.server_name,
&tokens.url,
OAuthCredentialsStoreMode::Auto,
AuthKeyringBackendKind::Direct,
)?,
None
);
assert!(
resolved
.store
.save(&store, &tokens.server_name, &tokens)
.is_err()
);
Ok::<_, anyhow::Error>(())
})
.join()
.unwrap()
})?;
let mut expected = vec![
expected_metric(
/*count*/ 1,
&[
("operation", "save"),
("actual_store", "file"),
("secure_outcome", "error"),
("fallback_reason", "secure_error"),
("secure_error", "unsupported"),
("originator", "none"),
],
),
expected_metric(/*count*/ 1, &[]),
expected_metric(
/*count*/ 1,
&[
("storage_phase", "pinned"),
("secure_outcome", "not_found"),
("outcome", "not_found"),
],
),
expected_metric(
/*count*/ 1,
&[
("storage_phase", "pinned"),
("operation", "save"),
("secure_outcome", "error"),
("outcome", "error"),
("secure_error", "timeout"),
],
),
];
expected.sort();
assert_eq!(operation_metrics(&metrics)?, expected);
metrics.shutdown()?;
Ok(())
}
fn isolated_process() -> Result<bool> {
if std::env::var_os("CODEX_AUTH_STORAGE_TEST_CHILD").is_none() {
let output = std::process::Command::new(std::env::current_exe()?)
.args([
"--exact",
std::thread::current().name().unwrap(),
"--nocapture",
])
.env("CODEX_AUTH_STORAGE_TEST_CHILD", "1")
.output()?;
assert!(
output.status.success(),
"{}{}",
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr)
);
return Ok(true);
}
Ok(false)
}
fn operation_metrics(metrics: &codex_otel::MetricsClient) -> Result<Vec<MetricRecord>> {
let snapshot = metrics.snapshot()?;
let mut actual = Vec::new();
for metric in snapshot
.scope_metrics()
.flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics)
{
if matches!(
metric.name(),
"codex.auth_storage.operation" | "codex.auth_storage.refresh_persist"
) && let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data()
{
for point in sum.data_points() {
actual.push((
metric.name().to_string(),
point.value(),
point
.attributes()
.map(|a| (a.key.to_string(), a.value.to_string()))
.collect::<BTreeMap<_, _>>(),
));
}
}
}
actual.sort();
Ok(actual)
}
fn expected_metric(count: u64, tags: &[(&str, &str)]) -> MetricRecord {
let mut attributes = BTreeMap::from([
("credential_kind", "mcp"),
("store_mode", "auto"),
("selected_store", "direct_keyring"),
("actual_store", "direct_keyring"),
("operation", "load"),
("secure_outcome", "success"),
("outcome", "success"),
("fallback_reason", "none"),
("secure_error", "none"),
("storage_phase", "policy"),
("originator", "codex_vscode"),
]);
attributes.extend(tags.iter().copied());
(
if attributes["operation"] == "refresh_persist" {
"codex.auth_storage.refresh_persist"
} else {
"codex.auth_storage.operation"
}
.into(),
count,
attributes
.into_iter()
.map(|(k, v)| (k.into(), v.into()))
.collect(),
)
}
fn test_metrics() -> Result<codex_otel::MetricsClient> {
Ok(codex_otel::install_global_metrics(
codex_otel::MetricsClient::new(
codex_otel::MetricsConfig::in_memory("test", "test", "1", Default::default())
.with_runtime_reader(),
)?,
))
}
#[test]
fn cleanup_failure_retains_policy_and_remains_best_effort() -> Result<()> {
if isolated_process()? {
return Ok(());
}
let _env = TempCodexHome::new();
let metrics = test_metrics()?;
let tokens = sample_tokens();
let key = compute_store_key(&tokens.server_name, &tokens.url)?;
std::fs::write(fallback_file_path()?, "{")?;
let mut expected = Vec::new();
AuthStorageOriginator::from_client_name("codex_vscode").sync_scope(|| -> Result<()> {
for (mode, mode_tag) in [
(OAuthCredentialsStoreMode::Auto, "auto"),
(OAuthCredentialsStoreMode::Keyring, "keyring"),
] {
let store = MockKeyringStore::default();
save_oauth_tokens_with_keyring_and_cleanup_file(
&store,
mode,
AuthKeyringBackendKind::Direct,
&tokens.server_name,
&tokens,
)?;
assert_eq!(
store.saved_value(&key),
Some(serde_json::to_string(&tokens)?)
);
expected.push(expected_metric(
/*count*/ 1,
&[
("store_mode", mode_tag),
("operation", "cleanup"),
("actual_store", "file"),
("secure_outcome", "not_attempted"),
("outcome", "error"),
],
));
}
Ok(())
})?;
expected.sort();
assert_eq!(operation_metrics(&metrics)?, expected);
metrics.shutdown()?;
Ok(())
}
#[test]
fn lock_errors_retain_typed_categories_through_wrappers() -> Result<()> {
if isolated_process()? {
return Ok(());
}
let metrics = test_metrics()?;
let mut expected = Vec::new();
AuthStorageOriginator::from_client_name("codex_vscode").sync_scope(|| {
for (duration, category) in [
(Duration::ZERO, "locked"),
(Duration::from_secs(/*secs*/ 1), "timeout"),
] {
for wrapped in [false, true] {
let error = OAuthStoreLockFailure::Timeout {
store: OAuthStore::Secrets,
path: "private-lock-path".into(),
acquire_timeout: duration,
};
let error = if wrapped {
Error::from(OAuthKeyringLoadError::StoreLock(error)).context("private context")
} else {
Error::from(error)
};
let result: Result<Option<()>> = Err(error);
let mut observation = telemetry::policy(
OAuthCredentialsStoreMode::Auto,
AuthKeyringBackendKind::Secrets,
Operation::Load,
);
observation.record_load_attempt(Store::Secrets, &result);
telemetry::record_secure_error(
&mut observation,
result.as_ref().unwrap_err().as_ref(),
);
}
expected.push(expected_metric(
/*count*/ 2,
&[
("selected_store", "secrets"),
("actual_store", "secrets"),
("secure_outcome", "error"),
("outcome", "error"),
("secure_error", category),
],
));
}
});
expected.sort();
assert_eq!(operation_metrics(&metrics)?, expected);
metrics.shutdown()?;
Ok(())
}
#[tokio::test(flavor = "current_thread")]
async fn coordinated_refresh_persistence_records_success_and_failure() -> Result<()> {
use rmcp::transport::auth::CredentialStore;
if isolated_process()? {
return Ok(());
}
let _env = TempCodexHome::new();
let metrics = test_metrics()?;
for fail_save in [false, true] {
let keyring = MockKeyringStore::default();
let tokens = sample_tokens();
let key = compute_store_key(&tokens.server_name, &tokens.url)?;
keyring.save(KEYRING_SERVICE, &key, &serde_json::to_string(&tokens)?)?;
let resolved = AuthStorageOriginator::from_client_name("codex_vscode")
.sync_scope(|| {
resolve_oauth_tokens_from_store_policy(
&keyring,
&tokens.server_name,
&tokens.url,
OAuthCredentialsStoreMode::Auto,
AuthKeyringBackendKind::Direct,
)
})?
.unwrap();
let store = OAuthCredentialStore::new(
tokens,
resolved.store,
keyring.clone(),
/*oauth_config*/ None,
);
let _guard = store.acquire_transaction_guard().await?;
let mut credentials = store.load().await?.unwrap();
credentials
.token_response
.as_mut()
.unwrap()
.set_access_token(AccessToken::new("refreshed-access".into()));
if fail_save {
keyring.set_error(
&key,
KeyringError::PlatformFailure(
std::io::Error::from(std::io::ErrorKind::TimedOut).into(),
),
);
}
let _ = operation_metrics(&metrics)?;
assert_eq!(store.save(credentials).await.is_err(), fail_save);
let outcome = if fail_save { "error" } else { "success" };
let expected = ["save", "refresh_persist"].map(|operation| {
expected_metric(
/*count*/ 1,
&[
("operation", operation),
("storage_phase", "pinned"),
("secure_outcome", outcome),
("outcome", outcome),
("secure_error", if fail_save { "timeout" } else { "none" }),
],
)
});
assert_eq!(operation_metrics(&metrics)?, expected);
}
metrics.shutdown()?;
Ok(())
}
#[expect(
clippy::await_holding_invalid_type,
reason = "The legacy authorization manager serializes refresh through its Tokio mutex"
)]
#[tokio::test(flavor = "current_thread")]
async fn legacy_refresh_persistence_records_each_write_once() -> Result<()> {
use wiremock::Mock;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
if isolated_process()? {
return Ok(());
}
let metrics = test_metrics()?;
for proactive in [false, true] {
for fail_save in [false, true] {
let (_env, server, initial) = test_context().await?;
save_oauth_tokens_to_file(&initial)?;
let file = fallback_file_path()?;
Mock::given(method("POST"))
.and(path("/oauth/token"))
.respond_with(move |_: &wiremock::Request| {
if fail_save {
std::fs::write(&file, "{").unwrap();
}
ResponseTemplate::new(/*s*/ 200).set_body_json(serde_json::json!({
"access_token": "refreshed-access", "token_type": "Bearer",
"refresh_token": "rotated-refresh", "expires_in": 3600,
}))
})
.expect(1)
.mount(&server)
.await;
let manager = Arc::new(Mutex::new(authorization_manager_for(&initial).await?));
let authority = AuthStorageOriginator::from_client_name("codex_vscode")
.sync_scope(ResolvedOAuthCredentialStore::file);
let persistor = OAuthPersistor::new(
initial.server_name.clone(),
initial.url.clone(),
Arc::clone(&manager),
authority,
Some(initial),
/*oauth_config*/ None,
);
let result = if proactive {
persistor.refresh_if_needed().await
} else {
manager.lock().await.refresh_token().await?;
persistor.persist_if_needed().await
};
assert_eq!(result.is_err(), fail_save);
let mut expected = Vec::new();
for operation in if proactive {
vec!["load", "save", "refresh_persist"]
} else {
vec!["save", "refresh_persist"]
} {
expected.push(expected_metric(
/*count*/ 1,
&[
("operation", operation),
("store_mode", "file"),
("actual_store", "file"),
("storage_phase", "pinned"),
("secure_outcome", "not_attempted"),
(
"outcome",
if fail_save && operation != "load" {
"error"
} else {
"success"
},
),
],
));
}
expected.sort();
assert_eq!(operation_metrics(&metrics)?, expected);
if !fail_save {
// The proactive path installs the persisted response; the reactive path updates
// its snapshot. Neither should emit another milestone without another write.
persistor.persist_if_needed().await?;
assert_eq!(operation_metrics(&metrics)?, Vec::new());
}
server.verify().await;
}
}
metrics.shutdown()?;
Ok(())
}
@@ -11,6 +11,7 @@ use anyhow::anyhow;
use anyhow::bail;
use codex_config::McpServerOAuthConfig;
use codex_exec_server::HttpClient;
use codex_otel::auth_storage::AuthStorageOriginator;
use rmcp::transport::AuthorizationManager;
use rmcp::transport::AuthorizationSession;
use rmcp::transport::auth::AuthorizationMetadata;
@@ -846,14 +847,15 @@ impl OauthLoginFlow {
let server_name = self.server_name.clone();
let (tx, rx) = oneshot::channel();
tokio::spawn(async move {
let task = async move {
let result = self.finish(/*emit_browser_url*/ false).await;
if let Err(err) = &result {
eprintln!("Failed to complete OAuth login for '{server_name}': {err:#}");
}
let _ = tx.send(result);
});
};
tokio::spawn(AuthStorageOriginator::current().scope(task));
rx
}
+17 -5
View File
@@ -21,6 +21,7 @@ use codex_config::types::AuthKeyringBackendKind;
use codex_config::types::McpServerEnvVar;
use codex_exec_server::HttpClient;
use codex_keyring_store::DefaultKeyringStore;
use codex_otel::auth_storage::AuthStorageOriginator;
use futures::FutureExt;
use futures::future::BoxFuture;
use http::HeaderMap;
@@ -168,6 +169,7 @@ enum TransportRecipe {
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
pinned_credential_store: Arc<OnceLock<ResolvedOAuthCredentialStore>>,
originator: AuthStorageOriginator,
http_client: Arc<dyn HttpClient>,
auth_provider: Option<SharedAuthProvider>,
redirect_mode: StreamableHttpRedirectMode,
@@ -608,6 +610,7 @@ impl RmcpClient {
store_mode,
keyring_backend_kind,
pinned_credential_store: Arc::new(OnceLock::new()),
originator: AuthStorageOriginator::current(),
http_client,
auth_provider,
redirect_mode,
@@ -1113,6 +1116,7 @@ impl RmcpClient {
store_mode,
keyring_backend_kind,
pinned_credential_store,
originator,
http_client,
auth_provider,
redirect_mode,
@@ -1148,13 +1152,18 @@ impl RmcpClient {
let oauth_keyring_backend_kind = *keyring_backend_kind;
let pinned_credential_store = Arc::clone(pinned_credential_store);
tokio::task::spawn_blocking(move || -> Result<Option<ResolvedOAuthTokens>> {
// Recovery can run after the startup scope ends, even when startup found
// no credentials and therefore could not pin a store with this attribution.
let originator = *originator;
let load = move || -> Result<Option<ResolvedOAuthTokens>> {
if let Some(store) = pinned_credential_store.get().copied() {
// Rebuilds reread the source selected during first construction. Only
// initial construction below evaluates configured store policy.
return store
.load(&DefaultKeyringStore, &oauth_server_name, &oauth_url)
.map(|tokens| tokens.map(|tokens| ResolvedOAuthTokens { tokens, store }));
.map(|tokens| {
tokens.map(|tokens| ResolvedOAuthTokens { tokens, store })
});
}
match resolve_oauth_tokens_from_store_policy(
@@ -1183,9 +1192,12 @@ impl RmcpClient {
Ok(None)
}
}
})
.await
.map_err(|error| anyhow!("OAuth credential loading task failed: {error}"))??
};
tokio::task::spawn_blocking(move || originator.sync_scope(load))
.await
.map_err(|error| {
anyhow!("OAuth credential loading task failed: {error}")
})??
} else {
None
};
@@ -1,5 +1,8 @@
mod streamable_http_test_support;
#[path = "streamable_http_recovery/telemetry_tests.rs"]
mod telemetry_tests;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
@@ -0,0 +1,152 @@
//! Checks attribution when a session rebuild discovers credentials after anonymous startup.
use super::streamable_http_test_support::arm_session_post_failure;
use super::streamable_http_test_support::call_echo_tool;
use super::streamable_http_test_support::expected_echo_result;
use super::streamable_http_test_support::initialize_client;
use super::streamable_http_test_support::spawn_streamable_http_server;
use codex_config::types::AuthKeyringBackendKind;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_exec_server::Environment;
use codex_otel::MetricsClient;
use codex_otel::MetricsConfig;
use codex_otel::auth_storage::AuthStorageOriginator;
use codex_rmcp_client::RmcpClient;
use codex_rmcp_client::StoredOAuthTokens;
use codex_rmcp_client::WrappedOAuthTokenResponse;
use codex_rmcp_client::save_oauth_tokens;
use oauth2::AccessToken;
use oauth2::basic::BasicTokenType;
use opentelemetry_sdk::metrics::data::AggregatedMetrics;
use opentelemetry_sdk::metrics::data::MetricData;
use pretty_assertions::assert_eq;
use rmcp::transport::auth::OAuthTokenResponse;
use std::collections::BTreeMap;
#[tokio::test]
async fn recovery_retains_originator_without_initial_credentials() -> anyhow::Result<()> {
if std::env::var_os("CODEX_MCP_RECOVERY_TELEMETRY_TEST_CHILD").is_none() {
let home = tempfile::tempdir()?;
let output = tokio::process::Command::new(std::env::current_exe()?)
.args([
"--exact",
"telemetry_tests::recovery_retains_originator_without_initial_credentials",
"--nocapture",
])
.env("CODEX_MCP_RECOVERY_TELEMETRY_TEST_CHILD", "1")
.env("CODEX_HOME", home.path())
.output()
.await?;
assert!(
output.status.success(),
"{}{}",
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr)
);
return Ok(());
}
let metrics = codex_otel::install_global_metrics(MetricsClient::new(
MetricsConfig::in_memory("test", "test", "1", Default::default()).with_runtime_reader(),
)?);
let (_server, base_url) = spawn_streamable_http_server().await?;
let url = format!("{base_url}/mcp");
let client = AuthStorageOriginator::from_client_name("codex_vscode")
.scope(async {
let client = RmcpClient::new_streamable_http_client(
"recovery-telemetry",
&url,
/*bearer_token*/ None,
/*http_headers*/ None,
/*env_http_headers*/ None,
OAuthCredentialsStoreMode::File,
AuthKeyringBackendKind::Direct,
Environment::default_for_tests().get_http_client(),
/*auth_provider*/ None,
)
.await?;
initialize_client(&client).await?;
anyhow::Ok(client)
})
.await?;
assert_eq!(
call_echo_tool(&client, "warmup").await?,
expected_echo_result("warmup")
);
for recovery in 0..3 {
if recovery == 1 {
// Credentials arrive after startup and the first anonymous recovery.
save_oauth_tokens(
"recovery-telemetry",
&StoredOAuthTokens {
server_name: "recovery-telemetry".into(),
url: url.clone(),
issuer: Some(url.clone()),
client_id: "client-id".into(),
token_response: WrappedOAuthTokenResponse(OAuthTokenResponse::new(
AccessToken::new("saved-access-token".into()),
BasicTokenType::Bearer,
Default::default(),
)),
expires_at: None,
},
OAuthCredentialsStoreMode::File,
AuthKeyringBackendKind::Direct,
)
.await?;
}
arm_session_post_failure(
&base_url,
/*status*/ 404,
/*remaining*/ 1,
/*www_authenticate_headers*/ &[],
)
.await?;
assert_eq!(
call_echo_tool(&client, "recovered").await?,
expected_echo_result("recovered")
);
}
let snapshot = metrics.snapshot()?;
let mut loads = Vec::new();
for metric in snapshot
.scope_metrics()
.flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics)
{
if metric.name() == "codex.auth_storage.operation"
&& let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data()
{
for point in sum.data_points() {
let tags: BTreeMap<_, _> = point
.attributes()
.map(|tag| (tag.key.to_string(), tag.value.to_string()))
.collect();
if tags["credential_kind"] == "mcp" && tags["operation"] == "load" {
loads.push((
tags["originator"].clone(),
tags["storage_phase"].clone(),
tags["outcome"].clone(),
point.value(),
));
}
}
}
}
loads.sort();
assert_eq!(
loads,
vec![
("codex_vscode".into(), "pinned".into(), "success".into(), 1),
(
"codex_vscode".into(),
"policy".into(),
"not_found".into(),
2
),
("codex_vscode".into(), "policy".into(), "success".into(), 1),
]
);
metrics.shutdown()?;
Ok(())
}