diff --git a/codex-rs/analytics/src/analytics_client_tests.rs b/codex-rs/analytics/src/analytics_client_tests.rs index bdd707be40..7cf59a2318 100644 --- a/codex-rs/analytics/src/analytics_client_tests.rs +++ b/codex-rs/analytics/src/analytics_client_tests.rs @@ -1,4 +1,6 @@ +use crate::client::AnalyticsEventsClient; use crate::client::AnalyticsEventsQueue; +use crate::client::AuthManagerRetention; use crate::events::AppServerRpcTransport; use crate::events::CodexAppMentionedEventRequest; use crate::events::CodexAppServerClientMetadata; @@ -85,6 +87,8 @@ use codex_app_server_protocol::TurnStatus as AppServerTurnStatus; use codex_app_server_protocol::TurnSteerParams; use codex_app_server_protocol::TurnSteerResponse; use codex_app_server_protocol::UserInput; +use codex_login::AuthManager; +use codex_login::CodexAuth; use codex_login::default_client::DEFAULT_ORIGINATOR; use codex_login::default_client::originator; use codex_plugin::AppConnectorId; @@ -795,6 +799,37 @@ fn app_used_dedupe_is_keyed_by_turn_and_connector() { assert_eq!(queue.should_enqueue_app_used(&turn_2, &app), true); } +#[tokio::test] +async fn default_client_retains_auth_manager() { + let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test")); + let weak_auth_manager = Arc::downgrade(&auth_manager); + + let _client = AnalyticsEventsClient::new( + auth_manager, + "http://localhost".to_string(), + None, + AuthManagerRetention::Strong, + ); + + assert!(weak_auth_manager.upgrade().is_some()); +} + +#[tokio::test] +async fn non_owning_client_does_not_retain_auth_manager() { + let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test")); + let weak_auth_manager = Arc::downgrade(&auth_manager); + + let _client = AnalyticsEventsClient::new( + Arc::clone(&auth_manager), + "http://localhost".to_string(), + None, + AuthManagerRetention::Weak, + ); + drop(auth_manager); + + assert!(weak_auth_manager.upgrade().is_none()); +} + #[test] fn thread_initialized_event_serializes_expected_shape() { let event = TrackEventRequest::ThreadInitialized(ThreadInitializedEvent { diff --git a/codex-rs/analytics/src/client.rs b/codex-rs/analytics/src/client.rs index 979260d22c..e836f466a0 100644 --- a/codex-rs/analytics/src/client.rs +++ b/codex-rs/analytics/src/client.rs @@ -34,6 +34,7 @@ use codex_plugin::PluginTelemetryMetadata; use std::collections::HashSet; use std::sync::Arc; use std::sync::Mutex; +use std::sync::Weak; use std::time::Duration; use tokio::sync::mpsc; @@ -54,10 +55,41 @@ pub struct AnalyticsEventsClient { analytics_enabled: Option, } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum AuthManagerRetention { + Strong, + Weak, +} + +enum RetainedAuthManager { + Strong(Arc), + Weak(Weak), +} + +impl RetainedAuthManager { + fn upgrade(&self) -> Option> { + match self { + Self::Strong(auth_manager) => Some(Arc::clone(auth_manager)), + Self::Weak(auth_manager) => auth_manager.upgrade(), + } + } +} + impl AnalyticsEventsQueue { - pub(crate) fn new(auth_manager: Arc, base_url: String) -> Self { + pub(crate) fn new( + auth_manager: Arc, + base_url: String, + retention: AuthManagerRetention, + ) -> Self { + let auth_manager = match retention { + AuthManagerRetention::Strong => RetainedAuthManager::Strong(auth_manager), + AuthManagerRetention::Weak => RetainedAuthManager::Weak(Arc::downgrade(&auth_manager)), + }; + Self::spawn(auth_manager, base_url) + } + + fn spawn(auth_manager: RetainedAuthManager, base_url: String) -> Self { let (sender, mut receiver) = mpsc::channel(ANALYTICS_EVENTS_QUEUE_SIZE); - let auth_manager = Arc::downgrade(&auth_manager); tokio::spawn(async move { let mut reducer = AnalyticsReducer::default(); while let Some(input) = receiver.recv().await { @@ -122,9 +154,10 @@ impl AnalyticsEventsClient { auth_manager: Arc, base_url: String, analytics_enabled: Option, + retention: AuthManagerRetention, ) -> Self { Self { - queue: AnalyticsEventsQueue::new(Arc::clone(&auth_manager), base_url), + queue: AnalyticsEventsQueue::new(auth_manager, base_url, retention), analytics_enabled, } } diff --git a/codex-rs/analytics/src/lib.rs b/codex-rs/analytics/src/lib.rs index ed0f1036ca..02bf88b2b2 100644 --- a/codex-rs/analytics/src/lib.rs +++ b/codex-rs/analytics/src/lib.rs @@ -7,6 +7,7 @@ use std::time::SystemTime; use std::time::UNIX_EPOCH; pub use client::AnalyticsEventsClient; +pub use client::AuthManagerRetention; pub use events::AppServerRpcTransport; pub use events::GuardianApprovalRequestSource; pub use events::GuardianReviewAnalyticsResult; diff --git a/codex-rs/app-server/src/analytics.rs b/codex-rs/app-server/src/analytics.rs deleted file mode 100644 index 24ed12d2ad..0000000000 --- a/codex-rs/app-server/src/analytics.rs +++ /dev/null @@ -1,16 +0,0 @@ -use std::sync::Arc; - -use codex_analytics::AnalyticsEventsClient; -use codex_core::config::Config; -use codex_login::AuthManager; - -pub(crate) fn analytics_events_client_from_config( - auth_manager: Arc, - config: &Config, -) -> AnalyticsEventsClient { - AnalyticsEventsClient::new( - auth_manager, - config.chatgpt_base_url.trim_end_matches('/').to_string(), - config.analytics_enabled, - ) -} diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 11451e180f..e712f0b7c5 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -3129,6 +3129,7 @@ mod tests { use anyhow::Result; use anyhow::anyhow; use anyhow::bail; + use codex_analytics::AuthManagerRetention; use codex_app_server_protocol::AutoReviewDecisionSource; use codex_app_server_protocol::GuardianApprovalReviewStatus; use codex_app_server_protocol::JSONRPCErrorError; @@ -3626,6 +3627,7 @@ mod tests { ), "http://localhost".to_string(), Some(false), + AuthManagerRetention::Strong, ), codex_home: codex_home.path().to_path_buf(), }; diff --git a/codex-rs/app-server/src/config_api.rs b/codex-rs/app-server/src/config_api.rs index e8bb82777c..5503f837fc 100644 --- a/codex-rs/app-server/src/config_api.rs +++ b/codex-rs/app-server/src/config_api.rs @@ -467,6 +467,7 @@ mod tests { use super::*; use crate::config_manager::apply_runtime_feature_enablement; use codex_analytics::AnalyticsEventsClient; + use codex_analytics::AuthManagerRetention; use codex_arg0::Arg0DispatchPaths; use codex_config::CloudRequirementsLoader; use codex_config::LoaderOverrides; @@ -832,6 +833,7 @@ mod tests { .trim_end_matches('/') .to_string(), analytics_config.analytics_enabled, + AuthManagerRetention::Strong, ), ); diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 144766cde0..13c634991c 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -50,7 +50,6 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; -use crate::analytics::analytics_events_client_from_config; use crate::config_manager::ConfigManager; use crate::error_code::INTERNAL_ERROR_CODE; use crate::error_code::INVALID_REQUEST_ERROR_CODE; @@ -67,7 +66,9 @@ use crate::transport::CHANNEL_CAPACITY; use crate::transport::ConnectionOrigin; use crate::transport::OutboundConnectionState; use crate::transport::route_outgoing_envelope; +use codex_analytics::AnalyticsEventsClient; use codex_analytics::AppServerRpcTransport; +use codex_analytics::AuthManagerRetention; use codex_app_server_protocol::ClientNotification; use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::ConfigWarningNotification; @@ -369,8 +370,15 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { let auth_manager = AuthManager::shared_from_config(args.config.as_ref(), args.enable_codex_api_key_env) .await; - let analytics_events_client = - analytics_events_client_from_config(Arc::clone(&auth_manager), args.config.as_ref()); + let analytics_events_client = AnalyticsEventsClient::new( + Arc::clone(&auth_manager), + args.config + .chatgpt_base_url + .trim_end_matches('/') + .to_string(), + args.config.analytics_enabled, + AuthManagerRetention::Weak, + ); let outgoing_message_sender = Arc::new(OutgoingMessageSender::new( outgoing_tx, analytics_events_client, diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 9aea97219a..6b63a0114e 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -1,7 +1,7 @@ #![deny(clippy::print_stdout, clippy::print_stderr)] -mod analytics; - +use codex_analytics::AnalyticsEventsClient; +use codex_analytics::AuthManagerRetention; use codex_arg0::Arg0DispatchPaths; use codex_config::ConfigLayerStackOrdering; use codex_config::LoaderOverrides; @@ -21,7 +21,6 @@ use std::sync::Arc; use std::sync::RwLock; use std::sync::atomic::AtomicBool; -use crate::analytics::analytics_events_client_from_config; use crate::config_manager::ConfigManager; use crate::message_processor::MessageProcessor; use crate::message_processor::MessageProcessorArgs; @@ -712,8 +711,12 @@ pub async fn run_main_with_transport_options( }); let processor_handle = tokio::spawn({ - let analytics_events_client = - analytics_events_client_from_config(Arc::clone(&auth_manager), &config); + let analytics_events_client = AnalyticsEventsClient::new( + Arc::clone(&auth_manager), + config.chatgpt_base_url.trim_end_matches('/').to_string(), + config.analytics_enabled, + AuthManagerRetention::Weak, + ); let outgoing_message_sender = Arc::new(OutgoingMessageSender::new( outgoing_tx, analytics_events_client, diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 5b568cc878..e9b473c9f3 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -25,6 +25,7 @@ use async_trait::async_trait; use axum::http::HeaderValue; use codex_analytics::AnalyticsEventsClient; use codex_analytics::AppServerRpcTransport; +use codex_analytics::AuthManagerRetention; use codex_app_server_protocol::AppListUpdatedNotification; use codex_app_server_protocol::AuthMode as LoginAuthMode; use codex_app_server_protocol::ChatgptAuthTokensRefreshParams; @@ -278,6 +279,7 @@ impl MessageProcessor { Arc::clone(&auth_manager), config.chatgpt_base_url.trim_end_matches('/').to_string(), config.analytics_enabled, + AuthManagerRetention::Strong, ); let thread_manager = Arc::new(ThreadManager::new( config.as_ref(), diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 1ebe2b4fc7..6a46b072b2 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -46,6 +46,7 @@ use async_channel::Sender; use chrono::Local; use chrono::Utc; use codex_analytics::AnalyticsEventsClient; +use codex_analytics::AuthManagerRetention; use codex_analytics::SubAgentThreadStartedInput; use codex_app_server_protocol::McpServerElicitationRequest; use codex_app_server_protocol::McpServerElicitationRequestParams; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index d4c0d337cf..4aa80e3f44 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -776,6 +776,7 @@ impl Session { Arc::clone(&auth_manager), config.chatgpt_base_url.trim_end_matches('/').to_string(), config.analytics_enabled, + AuthManagerRetention::Strong, ) }); let services = SessionServices { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 283220e8fa..0d1f45a059 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3423,6 +3423,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { Arc::clone(&auth_manager), config.chatgpt_base_url.trim_end_matches('/').to_string(), config.analytics_enabled, + AuthManagerRetention::Strong, ), hooks: Hooks::new(HooksConfig { legacy_notify_argv: config.notify.clone(), @@ -4782,6 +4783,7 @@ where Arc::clone(&auth_manager), config.chatgpt_base_url.trim_end_matches('/').to_string(), config.analytics_enabled, + AuthManagerRetention::Strong, ), hooks: Hooks::new(HooksConfig { legacy_notify_argv: config.notify.clone(),