From 8fef273391ab360f69fea20c2ae0f17443d9ea15 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Thu, 21 May 2026 14:09:51 -0700 Subject: [PATCH] codex: address PR review feedback (#23734) --- codex-rs/core/src/session/turn.rs | 15 ++- codex-rs/core/tests/suite/client.rs | 136 ++++++++++++++++++++++++++++ 2 files changed, 146 insertions(+), 5 deletions(-) diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 321a3e306f..8f81b794ad 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -70,6 +70,7 @@ use codex_async_utils::OrCancelExt; use codex_features::Feature; use codex_git_utils::get_git_repo_root; use codex_git_utils::get_git_repo_root_with_fs; +use codex_protocol::account::ProviderAccount; use codex_protocol::config_types::AutoCompactTokenLimitScope; use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::ServiceTier; @@ -599,11 +600,15 @@ async fn track_turn_resolved_config_analytics( let mut state = sess.state.lock().await; state.take_next_turn_is_first() }; - let plan_type = sess - .services - .auth_manager - .auth_cached() - .and_then(|auth| auth.account_plan_type()); + let plan_type = match turn_context + .provider + .account_state() + .ok() + .and_then(|state| state.account) + { + Some(ProviderAccount::Chatgpt { plan_type, .. }) => Some(plan_type), + Some(ProviderAccount::ApiKey) | Some(ProviderAccount::AmazonBedrock) | None => None, + }; sess.services .analytics_events_client .track_turn_resolved_config(TurnResolvedConfigFact { diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index f6731a0ea4..7e176d75db 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -74,6 +74,8 @@ use serde_json::json; use std::io::Write; use std::num::NonZeroU64; use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; use tempfile::TempDir; use toml::toml; use uuid::Uuid; @@ -3028,6 +3030,140 @@ fn create_dummy_codex_auth() -> CodexAuth { CodexAuth::create_dummy_chatgpt_auth_for_testing() } +async fn create_chatgpt_auth_with_plan( + plan_type: &str, + chatgpt_base_url: &str, +) -> anyhow::Result<(TempDir, CodexAuth)> { + let codex_home = TempDir::new()?; + write_auth_json( + &codex_home, + /*openai_api_key*/ None, + plan_type, + "Access-123", + Some("acc-123"), + ); + let auth = CodexAuth::from_auth_storage( + codex_home.path(), + AuthCredentialsStoreMode::File, + Some(chatgpt_base_url), + ) + .await? + .expect("auth.json should load ChatGPT auth"); + Ok((codex_home, auth)) +} + +async fn submit_hello_turn(codex: &codex_core::CodexThread) -> anyhow::Result<()> { + codex + .submit(Op::UserInput { + environments: None, + items: vec![UserInput::Text { + text: "hello".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + thread_settings: Default::default(), + }) + .await?; + wait_for_event(codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + Ok(()) +} + +async fn wait_for_turn_analytics_event(server: &MockServer) -> serde_json::Value { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + let requests = server.received_requests().await.unwrap_or_default(); + if let Some(event) = requests + .into_iter() + .filter(|request| request.url.path() == "/codex/analytics-events/events") + .find_map(|request| { + let payload: serde_json::Value = serde_json::from_slice(&request.body).ok()?; + payload["events"].as_array().and_then(|events| { + events + .iter() + .find(|event| event["event_type"] == "codex_turn_event") + .cloned() + }) + }) + { + return event; + } + if Instant::now() >= deadline { + panic!("timed out waiting for turn analytics request"); + } + tokio::time::sleep(Duration::from_millis(50)).await; + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn turn_analytics_uses_authenticated_provider_plan() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + let server = MockServer::start().await; + let _resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + + let chatgpt_base_url = server.uri(); + let (_auth_home, auth) = create_chatgpt_auth_with_plan("pro", &chatgpt_base_url).await?; + let mut builder = test_codex() + .with_auth(auth) + .with_config(move |config| config.chatgpt_base_url = chatgpt_base_url); + let codex = builder.build(&server).await?.codex; + + submit_hello_turn(&codex).await?; + + let event = wait_for_turn_analytics_event(&server).await; + assert_eq!(event["event_params"]["plan_type"], json!("pro")); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn turn_analytics_omits_plan_for_provider_without_openai_auth() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + let server = MockServer::start().await; + let _resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + + let provider = ModelProviderInfo { + name: "custom".to_string(), + base_url: Some(format!("{}/v1", server.uri())), + env_key: None, + experimental_bearer_token: None, + auth: None, + aws: None, + query_params: None, + env_key_instructions: None, + wire_api: WireApi::Responses, + http_headers: None, + env_http_headers: None, + request_max_retries: None, + stream_max_retries: None, + stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, + requires_openai_auth: false, + supports_websockets: false, + }; + let chatgpt_base_url = server.uri(); + let (_auth_home, auth) = create_chatgpt_auth_with_plan("pro", &chatgpt_base_url).await?; + let mut builder = test_codex().with_auth(auth).with_config(move |config| { + config.model_provider = provider; + config.model_provider_id = "custom".to_string(); + config.chatgpt_base_url = chatgpt_base_url; + }); + let codex = builder.build(&server).await?.codex; + + submit_hello_turn(&codex).await?; + + let event = wait_for_turn_analytics_event(&server).await; + assert_eq!(event["event_params"]["plan_type"], json!(null)); + Ok(()) +} + /// Scenario: /// - Turn 1: user sends U1; model streams deltas then a final assistant message A. /// - Turn 2: user sends U2; model streams a delta then the same final assistant message A.