diff --git a/codex-rs/analytics/src/analytics_client_tests.rs b/codex-rs/analytics/src/analytics_client_tests.rs index c6b83af5e0..164a31daf5 100644 --- a/codex-rs/analytics/src/analytics_client_tests.rs +++ b/codex-rs/analytics/src/analytics_client_tests.rs @@ -1,3 +1,4 @@ +use crate::client::AnalyticsEventsClient; use crate::client::AnalyticsEventsQueue; use crate::events::AppServerRpcTransport; use crate::events::CodexAcceptedLineFingerprintsEventParams; @@ -174,6 +175,7 @@ use std::collections::HashSet; use std::path::PathBuf; use std::sync::Arc; use std::sync::Mutex; +use std::time::SystemTime; use tokio::sync::mpsc; const TEST_PRODUCT_CLIENT_ID: &str = "codex_work_desktop"; @@ -1298,6 +1300,110 @@ index 1111111..2222222 assert!(event.event_params.line_fingerprints.is_empty()); } +#[tokio::test] +#[cfg(debug_assertions)] +async fn analytics_flush_delivers_completed_turn_with_file_diff() { + let nonce = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .expect("system clock should be after Unix epoch") + .as_nanos(); + let capture_path = std::env::temp_dir().join(format!( + "codex-analytics-turn-flush-{}-{nonce}.jsonl", + std::process::id() + )); + let auth_manager = codex_login::AuthManager::from_auth_for_testing( + codex_login::CodexAuth::create_dummy_chatgpt_auth_for_testing(), + ); + let client = AnalyticsEventsClient::new_for_capture_file(auth_manager, capture_path.clone()); + + for fact in [ + sample_initialize_fact(/*connection_id*/ 7), + AnalyticsFact::ClientResponse { + connection_id: 7, + request_id: RequestId::Integer(1), + response: Box::new(sample_thread_start_response( + "thread-2", /*ephemeral*/ false, "gpt-5", + )), + thread_originator: None, + }, + AnalyticsFact::ClientRequest { + connection_id: 7, + request_id: RequestId::Integer(3), + request: Box::new(sample_turn_start_request("thread-2", /*request_id*/ 3)), + }, + AnalyticsFact::ClientResponse { + connection_id: 7, + request_id: RequestId::Integer(3), + response: Box::new(sample_turn_start_response("turn-2")), + thread_originator: None, + }, + AnalyticsFact::Custom(CustomAnalyticsFact::TurnResolvedConfig(Box::new( + sample_turn_resolved_config("thread-2", "turn-2"), + ))), + AnalyticsFact::Notification(Box::new(sample_turn_started_notification( + "thread-2", "turn-2", + ))), + AnalyticsFact::Custom(CustomAnalyticsFact::TurnProfile(Box::new( + TurnProfileFact { + turn_id: "turn-2".to_string(), + profile: sample_turn_profile(), + }, + ))), + AnalyticsFact::Notification(Box::new(ServerNotification::TurnDiffUpdated( + TurnDiffUpdatedNotification { + thread_id: "thread-2".to_string(), + turn_id: "turn-2".to_string(), + diff: "\ +diff --git a/src/lib.rs b/src/lib.rs +index 1111111..2222222 +--- a/src/lib.rs ++++ b/src/lib.rs +@@ -0,0 +1 @@ ++let value = 1; +" + .to_string(), + }, + ))), + AnalyticsFact::Notification(Box::new(sample_turn_completed_notification( + "thread-2", + "turn-2", + AppServerTurnStatus::Completed, + /*codex_error_info*/ None, + ))), + ] { + client.record_fact(fact); + } + + client.flush().await; + + let contents = std::fs::read_to_string(&capture_path).expect("read captured analytics events"); + let event_types = contents + .lines() + .flat_map(|line| { + serde_json::from_str::(line) + .expect("parse captured analytics events")["events"] + .as_array() + .expect("captured events should be an array") + .iter() + .map(|event| { + event["event_type"] + .as_str() + .expect("captured event type should be a string") + .to_string() + }) + .collect::>() + }) + .collect::>(); + assert!(event_types.iter().any(|event| event == "codex_turn_event")); + assert!( + event_types + .iter() + .any(|event| event == "codex_accepted_line_fingerprints") + ); + + std::fs::remove_file(capture_path).expect("remove analytics capture file"); +} + #[test] fn compaction_event_serializes_expected_shape() { let event = TrackEventRequest::Compaction(Box::new(CodexCompactionEventRequest { diff --git a/codex-rs/analytics/src/client.rs b/codex-rs/analytics/src/client.rs index 853a15ce28..74ff18508a 100644 --- a/codex-rs/analytics/src/client.rs +++ b/codex-rs/analytics/src/client.rs @@ -50,14 +50,22 @@ use std::sync::Arc; use std::sync::Mutex; use std::time::Duration; use tokio::sync::mpsc; +use tokio::sync::oneshot; const ANALYTICS_EVENTS_QUEUE_SIZE: usize = 256; const ANALYTICS_EVENTS_TIMEOUT: Duration = Duration::from_secs(10); +// Covers two sequential POSTs plus queue/barrier scheduling; additional queued sends remain best-effort. +const ANALYTICS_EVENTS_FLUSH_TIMEOUT: Duration = Duration::from_secs(25); const ANALYTICS_EVENT_DEDUPE_MAX_KEYS: usize = 4096; +pub(crate) enum AnalyticsEventsQueueMessage { + Fact(Box), + Flush(oneshot::Sender<()>), +} + #[derive(Clone)] pub(crate) struct AnalyticsEventsQueue { - pub(crate) sender: mpsc::Sender, + pub(crate) sender: mpsc::Sender, pub(crate) app_used_emitted_keys: Arc>>, pub(crate) plugin_used_emitted_keys: Arc>>, } @@ -128,6 +136,13 @@ impl AnalyticsEventsQueue { tokio::spawn(async move { let mut reducer = AnalyticsReducer::default(); while let Some(input) = receiver.recv().await { + let input = match input { + AnalyticsEventsQueueMessage::Fact(input) => *input, + AnalyticsEventsQueueMessage::Flush(done_tx) => { + let _ = done_tx.send(()); + continue; + } + }; let mut events = Vec::new(); reducer.ingest(input, &mut events).await; send_track_events(&auth_manager, &destination, events).await; @@ -141,7 +156,11 @@ impl AnalyticsEventsQueue { } fn try_send(&self, input: AnalyticsFact) { - if self.sender.try_send(input).is_err() { + if self + .sender + .try_send(AnalyticsEventsQueueMessage::Fact(Box::new(input))) + .is_err() + { //TODO: add a metric for this tracing::warn!("dropping analytics events: queue is full"); } @@ -206,6 +225,29 @@ impl AnalyticsEventsClient { Self { queue: None } } + pub async fn flush(&self) { + let Some(queue) = self.queue.as_ref() else { + return; + }; + let (done_tx, done_rx) = oneshot::channel(); + let flushed = tokio::time::timeout(ANALYTICS_EVENTS_FLUSH_TIMEOUT, async { + if queue + .sender + .send(AnalyticsEventsQueueMessage::Flush(done_tx)) + .await + .is_err() + { + return false; + } + done_rx.await.is_ok() + }) + .await; + + if !matches!(flushed, Ok(true)) { + tracing::warn!("timed out or failed while flushing analytics events"); + } + } + pub fn track_skill_invocations( &self, tracking: TrackEventsContext, diff --git a/codex-rs/analytics/src/client_tests.rs b/codex-rs/analytics/src/client_tests.rs index 443167820c..41ff883b50 100644 --- a/codex-rs/analytics/src/client_tests.rs +++ b/codex-rs/analytics/src/client_tests.rs @@ -1,6 +1,7 @@ use super::AnalyticsEventsClient; use super::AnalyticsEventsDestination; use super::AnalyticsEventsQueue; +use super::AnalyticsEventsQueueMessage; #[cfg(debug_assertions)] use super::capture_track_events_request; #[cfg(debug_assertions)] @@ -57,6 +58,8 @@ use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::TurnStatus as AppServerTurnStatus; use codex_app_server_protocol::TurnSteerParams; use codex_app_server_protocol::TurnSteerResponse; +#[cfg(debug_assertions)] +use codex_login::AuthManager; use codex_utils_absolute_path::test_support::PathBufExt; use codex_utils_absolute_path::test_support::test_path_buf; use std::collections::HashSet; @@ -71,6 +74,18 @@ use std::time::SystemTime; use tokio::sync::mpsc; use tokio::sync::mpsc::error::TryRecvError; +#[cfg(debug_assertions)] +impl AnalyticsEventsClient { + pub(crate) fn new_for_capture_file(auth_manager: Arc, path: PathBuf) -> Self { + Self { + queue: Some(AnalyticsEventsQueue::new( + auth_manager, + AnalyticsEventsDestination::CaptureFile { path }, + )), + } + } +} + fn sample_accepted_line_fingerprint_event(thread_id: &str) -> TrackEventRequest { TrackEventRequest::AcceptedLineFingerprints(Box::new( CodexAcceptedLineFingerprintsEventRequest { @@ -197,7 +212,10 @@ fn unique_capture_path(name: &str) -> PathBuf { )) } -fn client_with_receiver() -> (AnalyticsEventsClient, mpsc::Receiver) { +fn client_with_receiver() -> ( + AnalyticsEventsClient, + mpsc::Receiver, +) { let (sender, receiver) = mpsc::channel(8); let queue = AnalyticsEventsQueue { sender, @@ -569,7 +587,8 @@ fn track_request_only_enqueues_analytics_relevant_requests() { client.track_request(/*connection_id*/ 7, request_id, &request); assert!(matches!( receiver.try_recv(), - Ok(AnalyticsFact::ClientRequest { .. }) + Ok(AnalyticsEventsQueueMessage::Fact(input)) + if matches!(*input, AnalyticsFact::ClientRequest { .. }) )); } @@ -596,7 +615,8 @@ fn track_response_only_enqueues_analytics_relevant_responses() { client.track_response(/*connection_id*/ 7, request_id, response); assert!(matches!( receiver.try_recv(), - Ok(AnalyticsFact::ClientResponse { .. }) + Ok(AnalyticsEventsQueueMessage::Fact(input)) + if matches!(*input, AnalyticsFact::ClientResponse { .. }) )); } @@ -608,6 +628,36 @@ fn track_response_only_enqueues_analytics_relevant_responses() { assert!(matches!(receiver.try_recv(), Err(TryRecvError::Empty))); } +#[tokio::test] +async fn flush_waits_for_preceding_fact_delivery() { + let (client, mut receiver) = client_with_receiver(); + client.track_request( + /*connection_id*/ 7, + RequestId::Integer(1), + &sample_turn_start_request(), + ); + + let flush = tokio::spawn(async move { client.flush().await }); + assert!(matches!( + receiver.recv().await, + Some(AnalyticsEventsQueueMessage::Fact(input)) + if matches!(*input, AnalyticsFact::ClientRequest { .. }) + )); + let done_tx = match receiver.recv().await { + Some(AnalyticsEventsQueueMessage::Flush(done_tx)) => done_tx, + _ => panic!("expected analytics flush barrier"), + }; + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + assert!(!flush.is_finished()); + done_tx.send(()).expect("flush receiver should remain open"); + flush.await.expect("flush task should complete"); +} + +#[tokio::test] +async fn flush_is_noop_when_analytics_is_disabled() { + AnalyticsEventsClient::disabled().flush().await; +} + #[test] fn track_event_request_batches_only_isolates_accepted_line_fingerprint_events() { let batches = track_event_request_batches(vec![ diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index c25fbe57fb..abc844167b 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -84,6 +84,8 @@ pub mod legacy_core { } const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); +// Covers the embedded drain, its analytics flush, and final task join. +const IN_PROCESS_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(45); /// Raw app-server request result for typed in-process requests. /// @@ -747,7 +749,7 @@ impl InProcessAppServerClient { .send(ClientCommand::Shutdown { response_tx }) .await .is_ok() - && let Ok(command_result) = timeout(SHUTDOWN_TIMEOUT, response_rx).await + && let Ok(command_result) = timeout(IN_PROCESS_SHUTDOWN_TIMEOUT, response_rx).await { command_result.map_err(|_| { IoError::new( @@ -757,7 +759,7 @@ impl InProcessAppServerClient { })??; } - if let Err(_elapsed) = timeout(SHUTDOWN_TIMEOUT, &mut worker_handle).await { + if let Err(_elapsed) = timeout(IN_PROCESS_SHUTDOWN_TIMEOUT, &mut worker_handle).await { worker_handle.abort(); let _ = worker_handle.await; } @@ -2284,4 +2286,32 @@ mod tests { .expect("shutdown should not wait for the 5s fallback timeout") .expect("shutdown should complete"); } + + #[tokio::test(start_paused = true)] + async fn shutdown_waits_for_in_process_drain() { + use std::sync::atomic::AtomicBool; + use std::sync::atomic::Ordering; + + let (command_tx, mut command_rx) = mpsc::channel(1); + let (_event_tx, event_rx) = mpsc::channel(1); + let completed = Arc::new(AtomicBool::new(false)); + let worker_completed = Arc::clone(&completed); + let worker_handle = tokio::spawn(async move { + let response_tx = match command_rx.recv().await { + Some(ClientCommand::Shutdown { response_tx }) => response_tx, + _ => panic!("expected shutdown command"), + }; + tokio::time::sleep(Duration::from_secs(30)).await; + worker_completed.store(true, Ordering::Release); + let _ = response_tx.send(Ok(())); + }); + let client = InProcessAppServerClient { + command_tx, + event_rx, + worker_handle, + }; + + client.shutdown().await.expect("shutdown should complete"); + assert!(completed.load(Ordering::Acquire)); + } } diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 84fec13b22..7be5c055fe 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -97,6 +97,8 @@ use tracing::warn; const IN_PROCESS_CONNECTION_ID: ConnectionId = ConnectionId(0); const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); +// Covers both bounded runtime drains plus the analytics client's 25-second best-effort flush. +const SHUTDOWN_ACK_TIMEOUT: Duration = Duration::from_secs(35); /// Default bounded channel capacity for in-process runtime queues. pub const DEFAULT_IN_PROCESS_CHANNEL_CAPACITY: usize = CHANNEL_CAPACITY; @@ -330,7 +332,7 @@ impl InProcessClientHandle { .await .is_ok() { - let _ = timeout(SHUTDOWN_TIMEOUT, done_rx).await; + let _ = timeout(SHUTDOWN_ACK_TIMEOUT, done_rx).await; } if let Err(_elapsed) = timeout(SHUTDOWN_TIMEOUT, &mut runtime_handle).await { @@ -394,6 +396,7 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult IoResult done_tx, + _ => panic!("expected in-process shutdown request"), + }; + tokio::time::sleep(SHUTDOWN_TIMEOUT + SHUTDOWN_TIMEOUT + Duration::from_secs(24)).await; + runtime_completed.store(true, Ordering::Release); + let _ = done_tx.send(()); + }); + let client = InProcessClientHandle { + client: InProcessClientSender { client_tx }, + event_rx, + runtime_handle, + _test_codex_home: None, + }; + + client + .shutdown() + .await + .expect("in-process runtime should shutdown cleanly"); + assert!(completed.load(Ordering::Acquire)); + } + #[test] fn guaranteed_delivery_helpers_cover_terminal_server_notifications() { assert!(server_notification_requires_delivery( diff --git a/codex-rs/exec/tests/suite/apply_patch.rs b/codex-rs/exec/tests/suite/apply_patch.rs index 808968871c..9683efbd95 100644 --- a/codex-rs/exec/tests/suite/apply_patch.rs +++ b/codex-rs/exec/tests/suite/apply_patch.rs @@ -145,3 +145,168 @@ async fn test_apply_patch_freeform_tool() -> anyhow::Result<()> { ); Ok(()) } + +#[cfg(not(target_os = "windows"))] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn shutdown_flushes_completed_turn_and_file_diff() -> anyhow::Result<()> { + use codex_login::AuthDotJson; + use codex_login::AuthKeyringBackendKind; + use codex_login::TokenData; + use codex_login::save_auth; + use codex_login::token_data::IdTokenInfo; + use codex_protocol::auth::AuthMode; + use core_test_support::skip_if_no_network; + use core_test_support::test_codex_exec::test_codex_exec; + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use std::time::Duration; + use wiremock::Mock; + use wiremock::ResponseTemplate; + use wiremock::matchers::method; + + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let run_dir = test.cwd_path(); + assert!( + std::process::Command::new("git") + .args(["init", "--quiet"]) + .current_dir(run_dir) + .status() + .context("initialize test git repository")? + .success() + ); + assert!( + std::process::Command::new("git") + .args([ + "remote", + "add", + "origin", + "https://github.com/test-org/turn-profile-test.git", + ]) + .current_dir(run_dir) + .status() + .context("add test git remote")? + .success() + ); + save_auth( + test.home_path(), + &AuthDotJson { + auth_mode: Some(AuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(TokenData { + id_token: IdTokenInfo { + raw_jwt: "header.eyJhbGciOiJub25lIn0.eyJzdWIiOiJ1c2VyLTEyMyJ9.".to_string(), + ..Default::default() + }, + access_token: "test-access-token".to_string(), + refresh_token: "test-refresh-token".to_string(), + account_id: None, + }), + last_refresh: None, + agent_identity: None, + personal_access_token: None, + bedrock_api_key: None, + }, + codex_login::AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + + let server = start_mock_server().await; + let deliveries = Arc::new(AtomicUsize::new(0)); + let completed_deliveries = Arc::clone(&deliveries); + Mock::given(method("POST")) + .and(|request: &wiremock::Request| request.url.path().ends_with("events")) + .respond_with(move |_request: &wiremock::Request| { + // A request is counted only once its delayed handler has completed. + std::thread::sleep(Duration::from_millis(750)); + completed_deliveries.fetch_add(1, Ordering::Release); + ResponseTemplate::new(200) + }) + .expect(3..) + .mount(&server) + .await; + mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_apply_patch_custom_tool_call( + "request_0", + "*** Begin Patch\n*** Add File: flushed.md\n+persist this line\n*** End Patch", + ), + ev_completed("request_0"), + ]), + sse(vec![ev_completed("request_1")]), + ], + ) + .await; + + test.cmd_with_server(&server) + .env_remove("CODEX_API_KEY") + .env_remove("CODEX_ANALYTICS_EVENTS_CAPTURE_FILE") + .arg("--skip-git-repo-check") + .arg("-s") + .arg("danger-full-access") + .arg("-c") + .arg("model_provider=\"test\"") + .arg("-c") + .arg(format!( + "model_providers.test={{name=\"test\",base_url={:?},wire_api=\"responses\",requires_openai_auth=false,supports_websockets=false}}", + format!("{}/v1", server.uri()) + )) + .arg("-c") + .arg("analytics.enabled=true") + .arg("-c") + .arg(format!("chatgpt_base_url={}", server.uri())) + .arg("write the file") + .assert() + .success(); + + let persisted = std::fs::read_to_string(run_dir.join("flushed.md")) + .context("read file written by apply_patch")?; + assert_eq!(persisted, "persist this line\n"); + + let requests = server + .received_requests() + .await + .context("read analytics requests")?; + let request_paths = requests + .iter() + .map(|request| request.url.path().to_string()) + .collect::>(); + let event_types = requests + .iter() + .filter(|request| request.url.path().ends_with("events")) + .flat_map(|request| { + serde_json::from_slice::(&request.body) + .ok() + .and_then(|body| { + body.get("events") + .and_then(serde_json::Value::as_array) + .cloned() + }) + .unwrap_or_default() + .into_iter() + .filter_map(|event| event["event_type"].as_str().map(str::to_owned)) + .collect::>() + }) + .collect::>(); + assert!( + deliveries.load(Ordering::Acquire) >= 3, + "the runtime exited before the delayed analytics deliveries completed (completed={}, events={event_types:?}, paths={request_paths:?})", + deliveries.load(Ordering::Acquire) + ); + assert!( + event_types.iter().any(|event| event == "codex_turn_event"), + "no completed event was delivered; request paths: {request_paths:?}" + ); + assert!( + event_types + .iter() + .any(|event| event == "codex_accepted_line_fingerprints"), + "no accepted-line event was delivered; request paths: {request_paths:?}" + ); + + Ok(()) +}