diff --git a/codex-rs/analytics/src/client.rs b/codex-rs/analytics/src/client.rs index a00c76da30..ef310acf2c 100644 --- a/codex-rs/analytics/src/client.rs +++ b/codex-rs/analytics/src/client.rs @@ -37,6 +37,10 @@ use crate::facts::TurnResolvedConfigFact; use crate::facts::TurnTokenUsageFact; use crate::guardian_v2::GuardianV2Event; use crate::now_unix_millis; +use crate::product_attribution::MAX_THREAD_PRODUCTS; +use crate::product_attribution::ThreadProductUpdate; +use crate::product_attribution::ThreadProducts; +use crate::product_attribution::product_event_batches; use crate::reducer::AnalyticsReducer; use crate::reducer::MAX_PLUGIN_MEASUREMENTS_PER_BATCH; use crate::reducer::tracked_tool_item_id; @@ -74,6 +78,8 @@ use std::collections::HashSet; use std::path::PathBuf; use std::sync::Arc; use std::sync::Mutex; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; use std::time::Duration; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -86,12 +92,18 @@ const ANALYTICS_EVENT_DEDUPE_MAX_KEYS: usize = 4096; pub(crate) enum AnalyticsEventsQueueMessage { Fact(Box), + ThreadProduct { + thread_id: String, + product_sku: Option, + generation: u64, + }, Flush(oneshot::Sender<()>), } #[derive(Clone)] pub(crate) struct AnalyticsEventsQueue { pub(crate) sender: mpsc::Sender, + pub(crate) product_generation: Arc, pub(crate) app_used_emitted_keys: Arc>>, pub(crate) plugin_used_emitted_keys: Arc>>, } @@ -159,26 +171,63 @@ fn analytics_capture_file_from_env() -> Option { impl AnalyticsEventsQueue { fn new(auth_manager: Arc, destination: AnalyticsEventsDestination) -> Self { let (sender, mut receiver) = mpsc::channel(ANALYTICS_EVENTS_QUEUE_SIZE); + let product_generation = Arc::new(AtomicU64::new(0)); + let worker_product_generation = Arc::clone(&product_generation); tokio::spawn(async move { let mut reducer = AnalyticsReducer::default(); + let mut thread_products = ThreadProducts::default(); + let mut registered_generation = 0; while let Some(input) = receiver.recv().await { + let generation = worker_product_generation.load(Ordering::Acquire); + if generation != registered_generation { + // A dropped registration invalidates both current and queued attribution. + thread_products = ThreadProducts::default(); + registered_generation = generation; + } let input = match input { AnalyticsEventsQueueMessage::Fact(input) => *input, + AnalyticsEventsQueueMessage::ThreadProduct { + thread_id, + product_sku, + generation, + } => { + if generation != registered_generation { + continue; + } + // Pending tool events belong to the configuration that collected them. + let mut events = Vec::new(); + if thread_products.get(&thread_id) != product_sku.as_deref() { + reducer.flush_thread(&thread_id, &mut events); + } + if product_sku.is_some() + && !thread_products.products.contains_key(&thread_id) + && thread_products.products.len() >= MAX_THREAD_PRODUCTS + && let Some(evicted) = thread_products.registered_threads.front() + { + reducer.flush_thread(evicted, &mut events); + } + send_track_events(&auth_manager, &destination, events, &thread_products) + .await; + thread_products.register(thread_id, product_sku); + continue; + } AnalyticsEventsQueueMessage::Flush(done_tx) => { let mut events = Vec::new(); reducer.flush(&mut events); - send_track_events(&auth_manager, &destination, events).await; + send_track_events(&auth_manager, &destination, events, &thread_products) + .await; let _ = done_tx.send(()); continue; } }; let mut events = Vec::new(); reducer.ingest(input, &mut events).await; - send_track_events(&auth_manager, &destination, events).await; + send_track_events(&auth_manager, &destination, events, &thread_products).await; } }); Self { sender, + product_generation, app_used_emitted_keys: Arc::new(Mutex::new(HashSet::new())), plugin_used_emitted_keys: Arc::new(Mutex::new(HashSet::new())), } @@ -238,6 +287,30 @@ impl AnalyticsEventsQueue { } impl AnalyticsEventsClient { + /// Registers this thread's product attribution without changing authentication. + /// A full queue clears shared attribution until threads register again, without blocking startup. + pub fn update_thread_product_sku(&self, thread_id: ThreadId, update: ThreadProductUpdate) { + let product_sku = match update { + ThreadProductUpdate::Set(product) => Some(product), + ThreadProductUpdate::Clear => None, + }; + let Some(queue) = &self.queue else { + return; + }; + if queue + .sender + .try_send(AnalyticsEventsQueueMessage::ThreadProduct { + thread_id: thread_id.to_string(), + product_sku, + generation: queue.product_generation.load(Ordering::Acquire), + }) + .is_err() + { + queue.product_generation.fetch_add(1, Ordering::AcqRel); + tracing::warn!("clearing analytics product attribution: queue is unavailable"); + } + } + pub fn new( auth_manager: Arc, base_url: String, @@ -861,6 +934,7 @@ async fn send_track_events( auth_manager: &AuthManager, destination: &AnalyticsEventsDestination, mut events: Vec, + thread_products: &ThreadProducts, ) { if events.is_empty() { return; @@ -879,8 +953,33 @@ async fn send_track_events( return; } - for events in track_event_request_batches(events) { - send_track_events_request(&auth, destination, events, &http_client_factory).await; + // Product streams are capped by product_event_batches. Parallel streams add no + // timeout waves to the existing sequential isolated-event request splitting. + let mut requests = tokio::task::JoinSet::new(); + for (product_sku, events) in product_event_batches(events, thread_products) { + let (auth, destination, factory) = ( + auth.clone(), + destination.clone(), + http_client_factory.clone(), + ); + let product_sku = product_sku.map(str::to_string); + requests.spawn(async move { + for events in track_event_request_batches(events) { + send_track_events_request( + &auth, + &destination, + events, + &factory, + product_sku.as_deref(), + ) + .await; + } + }); + } + while let Some(result) = requests.join_next().await { + if let Err(error) = result { + tracing::warn!(%error, "analytics product stream failed"); + } } } @@ -912,6 +1011,7 @@ async fn send_track_events_request( destination: &AnalyticsEventsDestination, events: Vec, http_client_factory: &codex_http_client::HttpClientFactory, + product_sku: Option<&str>, ) { if events.is_empty() { return; @@ -943,14 +1043,16 @@ async fn send_track_events_request( return; } }; - let response = client + let mut request = client .post(url) .timeout(ANALYTICS_EVENTS_TIMEOUT) .headers(codex_model_provider::auth_provider_from_auth(auth).to_auth_headers()) .header("Content-Type", "application/json") - .json(&payload) - .send() - .await; + .json(&payload); + if let Some(product_sku) = product_sku { + request = request.header("X-OpenAI-Product-Sku", product_sku); + } + let response = request.send().await; match response { Ok(response) if response.status().is_success() => {} @@ -974,6 +1076,9 @@ fn capture_track_events_request( return false; }; + use std::sync::PoisonError; + static CAPTURE_LOCK: Mutex<()> = Mutex::new(()); + let _capture_guard = CAPTURE_LOCK.lock().unwrap_or_else(PoisonError::into_inner); if let Err(err) = crate::analytics_capture::append_payload(path, payload) { tracing::error!( path = %path.display(), diff --git a/codex-rs/analytics/src/client_product_tests.rs b/codex-rs/analytics/src/client_product_tests.rs new file mode 100644 index 0000000000..029dea7fd1 --- /dev/null +++ b/codex-rs/analytics/src/client_product_tests.rs @@ -0,0 +1,345 @@ +//! Verify product attribution remains per thread across shared queues and HTTP batches. +//! Reconfiguration preserves buffered attribution; queue overflow falls back to unknown. + +use super::AnalyticsEventsClient; +use super::AnalyticsEventsDestination; +use super::AnalyticsEventsQueue; +use super::sample_regular_track_event; +use super::sample_thread_start_response; +use super::sample_turn_start_response; +use super::send_track_events; +use crate::ThreadProductUpdate; +use crate::events::AppServerRpcTransport; +use crate::facts::PluginMeasurementRow; +use crate::facts::PluginMeasurementsInput; +use crate::product_attribution::ThreadProducts; +use codex_app_server_protocol::ClientInfo; +use codex_app_server_protocol::ClientResponsePayload; +use codex_app_server_protocol::CollabAgentTool; +use codex_app_server_protocol::CollabAgentToolCallStatus; +use codex_app_server_protocol::InitializeParams; +use codex_app_server_protocol::ItemCompletedNotification; +use codex_app_server_protocol::ItemStartedNotification; +use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ServerNotification; +use codex_app_server_protocol::ThreadItem; +use codex_app_server_protocol::TurnStartedNotification; +use codex_login::AuthManager; +use codex_protocol::ThreadId; +use pretty_assertions::assert_eq; +use std::collections::BTreeMap; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; + +enum Responses { + Immediate, + PauseAfter(usize), + Together, +} + +type CapturedRequest = (Option, serde_json::Value); + +#[test] +fn product_registry_clears_unset_threads_and_bounds_retention() { + let mut products = ThreadProducts::default(); + products.register("unset".to_string(), Some("aeon".to_string())); + products.register("unset".to_string(), /*product*/ None); + for index in 0..=4096 { + products.register(format!("thread-{index}"), Some("tpp".to_string())); + } + let batches = crate::product_attribution::product_event_batches( + vec![ + sample_regular_track_event("unset"), + sample_regular_track_event("thread-0"), + sample_regular_track_event("thread-4096"), + ], + &products, + ); + assert_eq!( + batches + .into_iter() + .map(|(product, events)| (product, events.len())) + .collect::>(), + vec![(None, 2), (Some("tpp"), 1)] + ); + for index in 0..17 { + products.register(format!("thread-{index}"), Some(format!("product-{index}"))); + } + let events = (0..17) + .map(|index| sample_regular_track_event(&format!("thread-{index}"))) + .collect::>(); + let expected = serde_json::json!([[null, &events]]); + assert_eq!( + serde_json::to_value(crate::product_attribution::product_event_batches( + events, &products + )) + .unwrap(), + expected + ); +} + +async fn capture_requests( + count: usize, + responses: Responses, +) -> ( + AnalyticsEventsDestination, + tokio::task::JoinHandle>, + tokio::sync::oneshot::Receiver<()>, + tokio::sync::oneshot::Sender<()>, +) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let destination = AnalyticsEventsDestination::Http { + url: format!("http://{}/events", listener.local_addr().unwrap()), + }; + let (paused_tx, paused_rx) = tokio::sync::oneshot::channel(); + let (resume_tx, resume_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let mut pause = Some((paused_tx, resume_rx)); + let mut requests = Vec::new(); + let mut pending = Vec::new(); + for _ in 0..count { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + let (headers, body) = loop { + let mut buffer = [0; 4096]; + let count = stream.read(&mut buffer).await.unwrap(); + assert!(count > 0); + request.extend_from_slice(&buffer[..count]); + if let Some(end) = request.windows(4).position(|part| part == b"\r\n\r\n") { + let headers = String::from_utf8_lossy(&request[..end]).to_lowercase(); + let length: usize = headers + .lines() + .find_map(|line| line.strip_prefix("content-length: ")) + .unwrap() + .parse() + .unwrap(); + if request.len() >= end + 4 + length { + break ( + headers, + serde_json::from_slice(&request[end + 4..end + 4 + length]).unwrap(), + ); + } + } + }; + let product = headers + .lines() + .find_map(|line| line.strip_prefix("x-openai-product-sku: ")) + .map(str::to_string); + requests.push((product, body)); + if matches!(responses, Responses::PauseAfter(count) if count == requests.len()) { + let (paused, resume) = pause.take().unwrap(); + paused.send(()).unwrap(); + resume.await.unwrap(); + } + pending.push(stream); + if matches!(responses, Responses::Together) && requests.len() < count { + continue; + } + for mut stream in pending.drain(..) { + stream + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") + .await + .unwrap(); + } + } + requests + }); + (destination, server, paused_rx, resume_tx) +} + +#[tokio::test] +async fn product_streams_send_concurrently_and_preserve_event_order() { + // Withhold every response until all 16 requests arrive: serial delivery would time out. + let (destination, server, _, _) = capture_requests(/*count*/ 16, Responses::Together).await; + let manager = AuthManager::from_auth_for_testing( + codex_login::CodexAuth::create_dummy_chatgpt_auth_for_testing(), + ); + let mut products = ThreadProducts::default(); + for index in 0..16 { + products.register(format!("thread-{index}"), Some(format!("product-{index}"))); + } + let event = |index| { + let mut event = sample_regular_track_event(&format!("thread-{}", index % 16)); + if let crate::events::TrackEventRequest::SkillInvocation(skill) = &mut event { + skill.event_params.turn_id = Some(format!("turn-{index}")); + } + event + }; + let events = (0..32).map(event).collect(); + tokio::time::timeout( + std::time::Duration::from_secs(15), + send_track_events(&manager, &destination, events, &products), + ) + .await + .unwrap(); + let mut requests = server.await.unwrap(); + requests.sort_by(|a, b| a.0.cmp(&b.0)); + let mut expected = (0..16) + .map(|index| { + ( + Some(format!("product-{index}")), + serde_json::json!({"events": [event(index), event(index + 16)]}), + ) + }) + .collect::>(); + expected.sort_by(|a, b| a.0.cmp(&b.0)); + assert_eq!(requests, expected); +} + +#[tokio::test] +async fn buffered_tool_events_preserve_attribution_or_drop_it_on_queue_overflow() { + for (new_product, overflow) in [ + (Some("tpp"), false), + (None, false), + (Some("tpp"), true), + (None, true), + ] { + let (destination, server, paused, resume) = capture_requests( + if overflow { 5 } else { 2 }, + if overflow { + Responses::PauseAfter(2) + } else { + Responses::Immediate + }, + ) + .await; + let manager = AuthManager::from_auth_for_testing( + codex_login::CodexAuth::create_dummy_chatgpt_auth_for_testing(), + ); + let client = AnalyticsEventsClient { + queue: Some(AnalyticsEventsQueue::new(manager, destination)), + }; + let thread_id = ThreadId::new(); + client.update_thread_product_sku(thread_id, ThreadProductUpdate::Set("aeon".to_string())); + client.track_initialize( + /*connection_id*/ 1, + InitializeParams { + client_info: ClientInfo { + name: "test".to_string(), + title: None, + version: "1".to_string(), + }, + capabilities: None, + }, + "test_client".to_string(), + AppServerRpcTransport::InProcess, + ); + let mut response = sample_thread_start_response(); + let ClientResponsePayload::ThreadStart(started) = &mut response else { + unreachable!() + }; + started.thread.id = thread_id.to_string(); + client.track_response(/*connection_id*/ 1, RequestId::Integer(1), &response); + let ClientResponsePayload::TurnStart(started) = sample_turn_start_response() else { + unreachable!() + }; + client.track_notification(&ServerNotification::TurnStarted(TurnStartedNotification { + thread_id: thread_id.to_string(), + turn: started.turn, + })); + // Collaborator completions wait in the reducer for subsequent sampling evidence. + let completed = ItemCompletedNotification { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 2, + item: ThreadItem::CollabAgentToolCall { + id: "item".to_string(), + tool: CollabAgentTool::SendMessage, + status: CollabAgentToolCallStatus::Failed, + sender_thread_id: thread_id.to_string(), + receiver_thread_ids: Vec::new(), + prompt: None, + model: None, + reasoning_effort: None, + agents_states: Default::default(), + }, + }; + client.track_notification(&ServerNotification::ItemStarted(ItemStartedNotification { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + started_at_ms: 1, + item: completed.item.clone(), + })); + client.track_notification(&ServerNotification::ItemCompleted(completed)); + let measurement = PluginMeasurementsInput { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + item_id: "measurement".to_string(), + originator: "test_client".to_string(), + model_slug: None, + reasoning_effort: None, + plugin_id: "sites@test".to_string(), + execution_id: "execution".to_string(), + operation: "build".to_string(), + rows: vec![PluginMeasurementRow { + measurement_name: "duration_ms".to_string(), + number_value: 12.0, + dimensions: BTreeMap::new(), + }], + }; + if overflow { + client.track_plugin_measurements(measurement.clone()); + paused.await.unwrap(); + // Hold HTTP delivery while stale registrations fill the shared queue. + let capacity = client.queue.as_ref().unwrap().sender.capacity(); + for _ in 0..capacity { + client.update_thread_product_sku( + thread_id, + ThreadProductUpdate::Set("aeon".to_string()), + ); + } + } + client.update_thread_product_sku( + thread_id, + new_product.map_or(ThreadProductUpdate::Clear, |value| { + ThreadProductUpdate::Set(value.to_string()) + }), + ); + if overflow { + resume.send(()).unwrap(); + } + client.flush().await; + if overflow { + client.track_plugin_measurements(measurement.clone()); + client.flush().await; + client + .update_thread_product_sku(thread_id, ThreadProductUpdate::Set("tpp".to_string())); + client.track_plugin_measurements(measurement); + client.flush().await; + } + let requests = tokio::time::timeout(std::time::Duration::from_secs(5), server) + .await + .unwrap() + .unwrap(); + let expected = if overflow { + vec![ + (Some("aeon"), "codex_thread_initialized"), + (Some("aeon"), "codex_plugin_measurement_event"), + (None, "codex_collab_agent_tool_call_event"), + (None, "codex_plugin_measurement_event"), + (Some("tpp"), "codex_plugin_measurement_event"), + ] + } else { + vec![ + (Some("aeon"), "codex_thread_initialized"), + (Some("aeon"), "codex_collab_agent_tool_call_event"), + ] + } + .into_iter() + .map(|(product, event)| (product.map(str::to_string), event.to_string())) + .collect::>(); + assert_eq!( + requests + .into_iter() + .map(|(product, body)| ( + product, + body["events"][0]["event_type"] + .as_str() + .unwrap() + .to_string(), + )) + .collect::>(), + expected + ); + } +} diff --git a/codex-rs/analytics/src/client_tests.rs b/codex-rs/analytics/src/client_tests.rs index 9112070c9c..b4fdf1066b 100644 --- a/codex-rs/analytics/src/client_tests.rs +++ b/codex-rs/analytics/src/client_tests.rs @@ -4,7 +4,6 @@ use super::AnalyticsEventsQueue; use super::AnalyticsEventsQueueMessage; #[cfg(debug_assertions)] use super::capture_track_events_request; -#[cfg(debug_assertions)] use super::send_track_events; #[cfg(debug_assertions)] use super::send_track_events_request; @@ -60,6 +59,8 @@ use crate::facts::InvocationType; use crate::facts::PluginMeasurementRow; use crate::facts::PluginMeasurementsInput; use crate::facts::TrackEventsContext; +#[cfg(debug_assertions)] +use crate::product_attribution::ThreadProducts; use crate::reducer::MAX_PLUGIN_MEASUREMENTS_PER_BATCH; use codex_app_server_protocol::ApprovalsReviewer as AppServerApprovalsReviewer; use codex_app_server_protocol::AskForApproval as AppServerAskForApproval; @@ -108,6 +109,9 @@ use std::time::SystemTime; use tokio::sync::mpsc; use tokio::sync::mpsc::error::TryRecvError; +#[path = "client_product_tests.rs"] +mod product_tests; + #[cfg(debug_assertions)] impl AnalyticsEventsClient { pub(crate) fn new_for_capture_file(auth_manager: Arc, path: PathBuf) -> Self { @@ -289,6 +293,7 @@ fn client_with_receiver() -> ( let (sender, receiver) = mpsc::channel(8); let queue = AnalyticsEventsQueue { sender, + product_generation: Arc::new(std::sync::atomic::AtomicU64::new(0)), app_used_emitted_keys: Arc::new(Mutex::new(HashSet::new())), plugin_used_emitted_keys: Arc::new(Mutex::new(HashSet::new())), }; @@ -370,7 +375,14 @@ async fn capture_file_writes_exact_serialized_request() { let auth = codex_login::CodexAuth::create_dummy_chatgpt_auth_for_testing(); let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); - send_track_events_request(&auth, &destination, vec![event], &factory).await; + send_track_events_request( + &auth, + &destination, + vec![event], + &factory, + /*product_sku*/ None, + ) + .await; let contents = fs::read_to_string(&capture_path).expect("read capture file"); let lines = contents.lines().collect::>(); @@ -398,7 +410,14 @@ async fn capture_file_writes_final_batches_as_separate_lines() { for batch in track_event_request_batches(events) { let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); - send_track_events_request(&auth, &destination, batch, &factory).await; + send_track_events_request( + &auth, + &destination, + batch, + &factory, + /*product_sku*/ None, + ) + .await; } let contents = fs::read_to_string(&capture_path).expect("read capture file"); @@ -474,6 +493,7 @@ async fn api_key_auth_sends_only_plugin_events_to_codex_backend() { sample_artifact_operation_event("plugin-artifact"), plugin_measurement("plugin-measurement", "sample@test"), ], + &ThreadProducts::default(), ) .await; diff --git a/codex-rs/analytics/src/lib.rs b/codex-rs/analytics/src/lib.rs index f98e3ad9b4..6160dd704a 100644 --- a/codex-rs/analytics/src/lib.rs +++ b/codex-rs/analytics/src/lib.rs @@ -5,6 +5,7 @@ mod client; mod events; mod facts; mod guardian_v2; +mod product_attribution; mod reducer; mod thread_hint; @@ -78,6 +79,7 @@ pub use facts::TurnTokenUsageFact; pub use facts::build_track_events_context; pub use guardian_v2::GuardianV2Event; pub use guardian_v2::GuardianV2EventKind; +pub use product_attribution::ThreadProductUpdate; pub use thread_hint::ThreadHintStatus; pub use thread_hint::ThreadHintStatusEvent; diff --git a/codex-rs/analytics/src/product_attribution.rs b/codex-rs/analytics/src/product_attribution.rs new file mode 100644 index 0000000000..4fa872102e --- /dev/null +++ b/codex-rs/analytics/src/product_attribution.rs @@ -0,0 +1,109 @@ +//! Keep per-thread product attribution separate from the shared analytics identity. + +use crate::events::TrackEventRequest; +use std::collections::HashMap; +use std::collections::VecDeque; + +pub(crate) const MAX_THREAD_PRODUCTS: usize = 4096; +const MAX_PRODUCT_BATCHES: usize = 16; + +/// An explicit change to a thread's product attribution, independent of authentication. +#[derive(Clone)] +pub enum ThreadProductUpdate { + Set(String), + Clear, +} + +#[derive(Default)] +pub(crate) struct ThreadProducts { + pub(crate) products: HashMap, + pub(crate) registered_threads: VecDeque, +} + +impl ThreadProducts { + pub(crate) fn get(&self, thread_id: &str) -> Option<&str> { + self.products.get(thread_id).map(String::as_str) + } + + pub(crate) fn register(&mut self, thread_id: String, product: Option) { + self.registered_threads.retain(|id| id != &thread_id); + self.products.remove(&thread_id); + let Some(product) = product else { + return; + }; + // Attribution is best-effort. Evicted threads have no product header. + if self.products.len() >= MAX_THREAD_PRODUCTS + && let Some(oldest) = self.registered_threads.pop_front() + { + self.products.remove(&oldest); + } + self.registered_threads.push_back(thread_id.clone()); + self.products.insert(thread_id, product); + } +} + +/// Preserve order within each product, with a bounded number of concurrent streams. +/// Too many products drops only attribution, preserving the original event order. +pub(crate) fn product_event_batches( + events: Vec, + thread_products: &ThreadProducts, +) -> Vec<(Option<&str>, Vec)> { + let products: std::collections::HashSet<_> = events + .iter() + .map(|event| event_thread_id(event).and_then(|id| thread_products.get(id))) + .collect(); + if products.len() > MAX_PRODUCT_BATCHES { + return vec![(None, events)]; + } + let mut batches: Vec<(Option<&str>, Vec)> = Vec::new(); + for event in events { + let product = event_thread_id(&event).and_then(|thread_id| thread_products.get(thread_id)); + if let Some((_, batch)) = batches.iter_mut().find(|(sku, _)| *sku == product) { + batch.push(event); + } else { + batches.push((product, vec![event])); + } + } + batches +} + +fn event_thread_id(event: &TrackEventRequest) -> Option<&str> { + match event { + TrackEventRequest::SkillInvocation(event) => event.event_params.thread_id.as_deref(), + TrackEventRequest::ThreadInitialized(event) => Some(&event.event_params.thread_id), + TrackEventRequest::ThreadArchive(event) => Some(&event.event_params.thread_id), + TrackEventRequest::GuardianReview(event) => { + Some(&event.event_params.guardian_review.thread_id) + } + TrackEventRequest::GuardianV2(event) => Some(&event.event_params.guardian_v2.thread_id), + TrackEventRequest::AppMentioned(event) => event.event_params.thread_id.as_deref(), + TrackEventRequest::AppUsed(event) => event.event_params.app.thread_id.as_deref(), + TrackEventRequest::HookRun(event) => event.event_params.thread_id.as_deref(), + TrackEventRequest::Compaction(event) => Some(&event.event_params.thread_id), + TrackEventRequest::Goal(event) => Some(&event.event_params.thread_id), + TrackEventRequest::ThreadHintStatus(event) => Some(&event.event_params.thread_id), + TrackEventRequest::TurnEvent(event) => Some(&event.event_params.thread_id), + TrackEventRequest::TurnSteer(event) => Some(&event.event_params.thread_id), + TrackEventRequest::ArtifactOperation(event) => Some(&event.event_params.thread_id), + TrackEventRequest::CommandExecution(event) => Some(&event.event_params.base.thread_id), + TrackEventRequest::PluginMeasurement(event) => Some(&event.event_params.thread_id), + TrackEventRequest::FileChange(event) => Some(&event.event_params.base.thread_id), + TrackEventRequest::McpToolCall(event) => Some(&event.event_params.base.thread_id), + TrackEventRequest::DynamicToolCall(event) => Some(&event.event_params.base.thread_id), + TrackEventRequest::ControlToolCall(event) => Some(&event.event_params.base.thread_id), + TrackEventRequest::CollabAgentToolCall(event) => Some(&event.event_params.base.thread_id), + TrackEventRequest::WebSearch(event) => Some(&event.event_params.base.thread_id), + TrackEventRequest::ImageGeneration(event) => Some(&event.event_params.base.thread_id), + TrackEventRequest::AcceptedLineFingerprints(event) => Some(&event.event_params.thread_id), + TrackEventRequest::ReviewEvent(event) => Some(&event.event_params.thread_id), + TrackEventRequest::PluginUsed(event) => event.event_params.thread_id.as_deref(), + TrackEventRequest::PluginInstallRequested(event) => Some(&event.event_params.thread_id), + TrackEventRequest::PluginInstalled(_) + | TrackEventRequest::PluginUninstalled(_) + | TrackEventRequest::PluginEnabled(_) + | TrackEventRequest::PluginDisabled(_) + | TrackEventRequest::PluginInstallFailed(_) + | TrackEventRequest::ExternalAgentConfigImportCompleted(_) + | TrackEventRequest::ExternalAgentConfigImportFailure(_) => None, + } +} diff --git a/codex-rs/analytics/src/reducer.rs b/codex-rs/analytics/src/reducer.rs index 956f155523..7640223e42 100644 --- a/codex-rs/analytics/src/reducer.rs +++ b/codex-rs/analytics/src/reducer.rs @@ -1155,6 +1155,14 @@ impl AnalyticsReducer { } } + pub(crate) fn flush_thread(&mut self, thread_id: &str, out: &mut Vec) { + for ((pending_thread_id, _), state) in &mut self.tool_response_states { + if pending_thread_id == thread_id { + out.extend(state.pending_tool_events.drain(..)); + } + } + } + fn ingest_initialize( &mut self, connection_id: u64, diff --git a/codex-rs/analytics/src/tests/suite/plugins_tests.rs b/codex-rs/analytics/src/tests/suite/plugins_tests.rs index 268a36f3ad..584f5d0c84 100644 --- a/codex-rs/analytics/src/tests/suite/plugins_tests.rs +++ b/codex-rs/analytics/src/tests/suite/plugins_tests.rs @@ -286,6 +286,7 @@ fn plugin_used_dedupe_is_keyed_by_turn_and_plugin() { let (sender, _receiver) = mpsc::channel(1); let queue = AnalyticsEventsQueue { sender, + product_generation: Arc::new(std::sync::atomic::AtomicU64::new(0)), app_used_emitted_keys: Arc::new(Mutex::new(HashSet::new())), plugin_used_emitted_keys: Arc::new(Mutex::new(HashSet::new())), }; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 77085befc4..824cf589cd 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -22,6 +22,7 @@ use crate::shell_snapshot::ShellSnapshot; use crate::shell_snapshot::SnapshotCredentialBrokerState; use crate::state::ActiveTurn; use crate::turn_metadata::ExecutionMetadata; +use codex_analytics::ThreadProductUpdate; use codex_attachment_store::AttachmentStore; use codex_extension_api::ExtensionDataInit; use codex_http_client::ClientRouteClass; @@ -1599,6 +1600,9 @@ impl Session { } let analytics_events_client = if config.analytics_enabled == Some(false) { + if let Some(client) = &analytics_events_client { + client.update_thread_product_sku(thread_id, ThreadProductUpdate::Clear); + } AnalyticsEventsClient::disabled() } else { analytics_events_client.unwrap_or_else(|| { @@ -1609,6 +1613,13 @@ impl Session { ) }) }; + analytics_events_client.update_thread_product_sku( + thread_id, + match &config.apps_mcp_product_sku { + Some(product) => ThreadProductUpdate::Set(product.clone()), + None => ThreadProductUpdate::Clear, + }, + ); for item in initial_history.get_rollout_items() { match item { RolloutItem::Compacted(compacted) => { diff --git a/codex-rs/core/tests/suite/plugins.rs b/codex-rs/core/tests/suite/plugins.rs index c817c988f0..2f13c1803a 100644 --- a/codex-rs/core/tests/suite/plugins.rs +++ b/codex-rs/core/tests/suite/plugins.rs @@ -5,6 +5,9 @@ use std::time::Duration; use std::time::Instant; use anyhow::Result; +use codex_analytics::AnalyticsEventsClient; +use codex_analytics::PluginMeasurementRow; +use codex_analytics::PluginMeasurementsInput; use codex_config::LoaderOverrides; use codex_core::TurnInputRequest; use codex_core::config::Config; @@ -24,6 +27,7 @@ use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::Settings; use codex_protocol::config_types::TrustLevel; +use codex_protocol::mcp::ClientMcpExtensions; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; @@ -363,6 +367,105 @@ fn searched_plugin_tools( ) } +#[tokio::test] +async fn shared_analytics_client_preserves_session_products() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + Mock::given(path("/codex/analytics-events/events")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server) + .await; + let client = AnalyticsEventsClient::new( + codex_core::test_support::auth_manager_from_auth( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + ), + server.uri(), + /*analytics_enabled*/ Some(true), + ); + let mut products = [Some("aeon"), Some("tpp"), None]; + let mut sessions = Vec::new(); + for product in products { + sessions.push( + test_codex() + .with_analytics_events_client(client.clone()) + .with_config(move |config| { + config.apps_mcp_product_sku = product.map(str::to_string); + }) + .build_with_auto_env(&server) + .await?, + ); + } + // All Sessions have registered before events are interleaved on their shared client. + let mut expected = Vec::new(); + for (position, index) in [0, 1, 2, 0, 0].into_iter().enumerate() { + if position == 4 { + client.flush().await; + let session = &mut sessions[0]; + session.codex.ensure_rollout_materialized().await; + session.codex.shutdown_and_wait().await?; + session + .thread_manager + .remove_thread(&session.session_configured.thread_id) + .await; + let mut config = session.config.clone(); + config.analytics_enabled = Some(false); + let resumed = session + .thread_manager + .resume_legacy_thread_from_rollout( + config, + session.codex.rollout_path().expect("rollout path"), + session.thread_manager.auth_manager(), + /*parent_trace*/ None, + ClientMcpExtensions::default(), + ) + .await?; + assert_eq!(resumed.thread_id, session.session_configured.thread_id); + session.codex = resumed.thread; + products[0] = None; + } + let thread_id = sessions[index].session_configured.thread_id.to_string(); + client.track_plugin_measurements(PluginMeasurementsInput { + thread_id: thread_id.clone(), + turn_id: "turn".into(), + item_id: "item".into(), + originator: "test_client".into(), + model_slug: None, + reasoning_effort: None, + plugin_id: "sample@test".into(), + execution_id: "execution".into(), + operation: "build".into(), + rows: vec![PluginMeasurementRow { + measurement_name: "duration_ms".into(), + number_value: 1.0, + dimensions: Default::default(), + }], + }); + expected.push((thread_id, products[index].map(str::to_string))); + } + client.flush().await; + let mut actual = Vec::new(); + for request in server.received_requests().await.unwrap_or_default() { + if request.url.path() != "/codex/analytics-events/events" { + continue; + } + let product = request + .headers + .get("x-openai-product-sku") + .map(|value| value.to_str().unwrap().to_string()); + let body: serde_json::Value = serde_json::from_slice(&request.body)?; + for event in body["events"].as_array().unwrap() { + if event["event_type"] == "codex_plugin_measurement_event" { + let thread_id = event["event_params"]["thread_id"].as_str().unwrap(); + actual.push((thread_id.to_string(), product.clone())); + } + } + } + actual.sort(); + expected.sort(); + assert_eq!(actual, expected); + Ok(()) +} + #[test_case(false; "classic shell")] #[test_case(true; "zsh-fork shell")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)]