Attribute analytics requests to each thread's product SKU (#49117)

## Why

Sessions can share an analytics client while using different product SKUs. Analytics requests need per-thread attribution to reflect each session's configuration.

## What changed

- Register each thread's `apps_mcp_product_sku` and send `X-OpenAI-Product-Sku` on analytics requests grouped by product, without changing authentication.
- Send product groups concurrently while preserving event order within each group.
- Flush buffered tool events before changing or evicting their product attribution, and clear attribution when a session disables analytics.
- Bound retained thread registrations and concurrent product groups. Drop attribution when registration cannot be queued or a batch exceeds the product-group limit.

## Testing

Add tests for shared-session attribution, clearing attribution on resume with analytics disabled, buffered events across product changes, queue overflow recovery, registry limits, and concurrent delivery with preserved event order.

GitOrigin-RevId: 25a98788cd44d4a85f97ee87aadf3a6909607f6a
This commit is contained in:
papayo-oai
2026-09-29 00:49:23 +00:00
committed by copyberry
parent 69043f05c4
commit 2e6cc4ed8d
9 changed files with 715 additions and 11 deletions
+113 -8
View File
@@ -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<AnalyticsFact>),
ThreadProduct {
thread_id: String,
product_sku: Option<String>,
generation: u64,
},
Flush(oneshot::Sender<()>),
}
#[derive(Clone)]
pub(crate) struct AnalyticsEventsQueue {
pub(crate) sender: mpsc::Sender<AnalyticsEventsQueueMessage>,
pub(crate) product_generation: Arc<AtomicU64>,
pub(crate) app_used_emitted_keys: Arc<Mutex<HashSet<(String, String)>>>,
pub(crate) plugin_used_emitted_keys: Arc<Mutex<HashSet<(String, String)>>>,
}
@@ -159,26 +171,63 @@ fn analytics_capture_file_from_env() -> Option<PathBuf> {
impl AnalyticsEventsQueue {
fn new(auth_manager: Arc<AuthManager>, 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<AuthManager>,
base_url: String,
@@ -861,6 +934,7 @@ async fn send_track_events(
auth_manager: &AuthManager,
destination: &AnalyticsEventsDestination,
mut events: Vec<TrackEventRequest>,
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<TrackEventRequest>,
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(),
@@ -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<String>, 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<_>>(),
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::<Vec<_>>();
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<Vec<CapturedRequest>>,
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::<Vec<_>>();
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::<Vec<_>>();
assert_eq!(
requests
.into_iter()
.map(|(product, body)| (
product,
body["events"][0]["event_type"]
.as_str()
.unwrap()
.to_string(),
))
.collect::<Vec<_>>(),
expected
);
}
}
+23 -3
View File
@@ -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<AuthManager>, 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::<Vec<_>>();
@@ -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;
+2
View File
@@ -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;
@@ -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<String, String>,
pub(crate) registered_threads: VecDeque<String>,
}
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<String>) {
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<TrackEventRequest>,
thread_products: &ThreadProducts,
) -> Vec<(Option<&str>, Vec<TrackEventRequest>)> {
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<TrackEventRequest>)> = 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,
}
}
+8
View File
@@ -1155,6 +1155,14 @@ impl AnalyticsReducer {
}
}
pub(crate) fn flush_thread(&mut self, thread_id: &str, out: &mut Vec<TrackEventRequest>) {
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,
@@ -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())),
};
+11
View File
@@ -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) => {
+103
View File
@@ -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)]