diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 234023d280..56a74c464d 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4389,6 +4389,7 @@ dependencies = [ "libc", "memchr", "oauth2", + "opentelemetry_sdk", "pretty_assertions", "rmcp", "serde", diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index ac2a55ba79..16adf8eb21 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -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(¶ms.server, ¶ms.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(()) } } diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index ad4fbf4828..65f7c5b445 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -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 { diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index f40ca17224..6d1d90b74c 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -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, diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 889ae89a2e..e7c95576db 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -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 diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index 800527900f..375f88c45f 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -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); } diff --git a/codex-rs/core/tests/suite/mcp_storage_telemetry_tests.rs b/codex-rs/core/tests/suite/mcp_storage_telemetry_tests.rs new file mode 100644 index 0000000000..71a9701150 --- /dev/null +++ b/codex-rs/core/tests/suite/mcp_storage_telemetry_tests.rs @@ -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> { + 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!["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::>(), + originators + ); + for thread in threads { + thread.shutdown_and_wait().await?; + } + fixture.codex.shutdown_and_wait().await?; + metrics.shutdown()?; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 2833288bb6..5400f0f4ef 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -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=="; diff --git a/codex-rs/rmcp-client/Cargo.toml b/codex-rs/rmcp-client/Cargo.toml index b133ddaf46..671bfd9ecc 100644 --- a/codex-rs/rmcp-client/Cargo.toml +++ b/codex-rs/rmcp-client/Cargo.toml @@ -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 diff --git a/codex-rs/rmcp-client/src/ema_identity_tests.rs b/codex-rs/rmcp-client/src/ema_identity_tests.rs index 53c7ef497a..56222f744d 100644 --- a/codex-rs/rmcp-client/src/ema_identity_tests.rs +++ b/codex-rs/rmcp-client/src/ema_identity_tests.rs @@ -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 }, diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index c4fa7a1f95..f3127d399a 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -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> { - 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( 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( 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( 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 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::().is_some() => Err(error), + Err(error) if error.downcast_ref::().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( server_name: &str, url: &str, ) -> Result { - 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( @@ -685,7 +717,7 @@ fn delete_oauth_tokens_from_direct_keyring( 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( @@ -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); diff --git a/codex-rs/rmcp-client/src/oauth/credential_store.rs b/codex-rs/rmcp-client/src/oauth/credential_store.rs index 2d79fe403b..c191e6936a 100644 --- a/codex-rs/rmcp-client/src/oauth/credential_store.rs +++ b/codex-rs/rmcp-client/src/oauth/credential_store.rs @@ -243,9 +243,11 @@ impl CredentialStore for OAuthCredentialStore .map_err(credential_store_error)?; let tokens = tokio::task::spawn_blocking(move || -> Result { 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 diff --git a/codex-rs/rmcp-client/src/oauth/ema_identity.rs b/codex-rs/rmcp-client/src/oauth/ema_identity.rs index 6f804a245f..bf4ea0d17a 100644 --- a/codex-rs/rmcp-client/src/oauth/ema_identity.rs +++ b/codex-rs/rmcp-client/src/oauth/ema_identity.rs @@ -45,7 +45,7 @@ impl StoredOAuthCredentialSnapshot { &self, keyring_store: &K, ) -> Result { - if !matches!(self.store, ResolvedOAuthCredentialStore::Keyring(_)) { + if self.store == ResolvedOAuthCredentialStore::file() { bail!("enterprise IdP credentials require keyring storage"); } let previous = &self.credentials; diff --git a/codex-rs/rmcp-client/src/oauth/ema_identity_tests.rs b/codex-rs/rmcp-client/src/oauth/ema_identity_tests.rs index c66721f395..bd804998d2 100644 --- a/codex-rs/rmcp-client/src/oauth/ema_identity_tests.rs +++ b/codex-rs/rmcp-client/src/oauth/ema_identity_tests.rs @@ -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": { diff --git a/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs b/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs index e5341ca73b..641f930c92 100644 --- a/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs +++ b/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs @@ -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()) diff --git a/codex-rs/rmcp-client/src/oauth/resolved_store.rs b/codex-rs/rmcp-client/src/oauth/resolved_store.rs index 904f5d9670..bdd29b58f9 100644 --- a/codex-rs/rmcp-client/src/oauth/resolved_store.rs +++ b/codex-rs/rmcp-client/src/oauth/resolved_store.rs @@ -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> { - 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> { - 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( + 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 { - 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 Result> { + 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 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 Result> { - 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::().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}")), }, } } diff --git a/codex-rs/rmcp-client/src/oauth/telemetry.rs b/codex-rs/rmcp-client/src/oauth/telemetry.rs new file mode 100644 index 0000000000..cbfa47aeb7 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/telemetry.rs @@ -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::() { + Some(OAuthKeyringLoadError::StoreLock(lock)) => Some(lock), + Some(OAuthKeyringLoadError::Backend(_)) | None => { + current.downcast_ref::() + } + }; + 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, + } +} diff --git a/codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs index 86882fc464..575ed835cd 100644 --- a/codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs +++ b/codex-rs/rmcp-client/src/oauth/tests/credential_store_tests.rs @@ -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, ); diff --git a/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs index 2951504e59..e938bcec48 100644 --- a/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs +++ b/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs @@ -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 { 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 { +pub(super) async fn authorization_manager_for( + tokens: &StoredOAuthTokens, +) -> Result { let oauth_http_client = Arc::new(OAuthHttpClientAdapter::new( Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( OutboundProxyPolicy::ReqwestDefault, diff --git a/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs index b8428f5e5e..8213c5ce42 100644 --- a/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs +++ b/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs @@ -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)); diff --git a/codex-rs/rmcp-client/src/oauth/tests/telemetry_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/telemetry_tests.rs new file mode 100644 index 0000000000..a3af6f2a47 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/tests/telemetry_tests.rs @@ -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); + +#[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 { + 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> { + 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::>(), + )); + } + } + } + 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 { + 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> = 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(()) +} diff --git a/codex-rs/rmcp-client/src/perform_oauth_login.rs b/codex-rs/rmcp-client/src/perform_oauth_login.rs index a271ed342e..a9053f6b7f 100644 --- a/codex-rs/rmcp-client/src/perform_oauth_login.rs +++ b/codex-rs/rmcp-client/src/perform_oauth_login.rs @@ -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 } diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index b8817d2eba..39dc2cdda4 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -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>, + originator: AuthStorageOriginator, http_client: Arc, auth_provider: Option, 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> { + // 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> { 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 }; diff --git a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs index 817587a9b7..8ac200fe03 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs @@ -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; diff --git a/codex-rs/rmcp-client/tests/streamable_http_recovery/telemetry_tests.rs b/codex-rs/rmcp-client/tests/streamable_http_recovery/telemetry_tests.rs new file mode 100644 index 0000000000..71171c5f73 --- /dev/null +++ b/codex-rs/rmcp-client/tests/streamable_http_recovery/telemetry_tests.rs @@ -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(()) +}