mirror of
https://github.com/openai/codex.git
synced 2026-11-08 20:40:33 +00:00
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:
@@ -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
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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())),
|
||||
};
|
||||
|
||||
@@ -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) => {
|
||||
|
||||
@@ -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)]
|
||||
|
||||
Reference in New Issue
Block a user