diff --git a/codex-rs/app-server/src/turn_cost_worker.rs b/codex-rs/app-server/src/turn_cost_worker.rs index 43b7b8e80e..b11efb9eef 100644 --- a/codex-rs/app-server/src/turn_cost_worker.rs +++ b/codex-rs/app-server/src/turn_cost_worker.rs @@ -13,6 +13,7 @@ use codex_protocol::auth::AuthMode; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use std::collections::HashMap; +use std::collections::HashSet; use std::sync::Arc; use std::time::Duration; use std::time::Instant; @@ -26,6 +27,7 @@ const REQUEST_TIMEOUT: Duration = Duration::from_secs(15); const OBSERVATION_CHANNEL_CAPACITY: usize = 16_384; const MAX_TRACKED_TURNS: usize = 4_096; const MAX_QUERY_TURNS: usize = 100; +const MAX_TRACKED_RESPONSES_PER_TURN: usize = 4_096; const MAX_STALLED_POLL_ATTEMPTS: u8 = 5; pub(crate) struct TurnCostWorker { @@ -45,7 +47,9 @@ enum TurnCostObservationKind { Started { session_telemetry: Box, }, - ResponseCompleted, + ResponseCompleted { + response_id: String, + }, Finished { interrupted: bool, }, @@ -67,7 +71,7 @@ enum TurnCostStatus { struct TurnCostEntry { thread_id: ThreadId, session_telemetry: SessionTelemetry, - expected_response_count: u64, + expected_response_ids: HashSet, status: TurnCostStatus, next_poll_at: Instant, attempt_count: u8, @@ -81,7 +85,7 @@ struct WorkerRuntime { #[derive(Clone)] enum TurnCostBackend { - OpenAiApiKey(Arc), + OpenAi(Arc), ModelProvider(SharedModelProvider), } @@ -110,7 +114,7 @@ impl TurnCostWorker { } let is_openai = config.model_provider.is_openai(); let backend = if is_openai { - TurnCostBackend::OpenAiApiKey(Arc::clone(&auth_manager)) + TurnCostBackend::OpenAi(Arc::clone(&auth_manager)) } else { TurnCostBackend::ModelProvider(create_model_provider( config.model_provider.clone(), @@ -165,11 +169,11 @@ impl TurnCostWorkerHandle { if thread_config.model_provider != self.config.model_provider { return; } - if let TurnCostBackend::OpenAiApiKey(auth_manager) = &self.backend { + if let TurnCostBackend::OpenAi(auth_manager) = &self.backend { let Some(auth) = auth_manager.auth_cached() else { return; }; - if !auth.is_api_key_auth() { + if !auth.is_api_key_auth() && !auth.is_chatgpt_auth() { return; } } @@ -177,7 +181,9 @@ impl TurnCostWorkerHandle { EventMsg::TurnStarted(_) => TurnCostObservationKind::Started { session_telemetry: Box::new(session_telemetry()), }, - EventMsg::RawResponseCompleted(_) => TurnCostObservationKind::ResponseCompleted, + EventMsg::RawResponseCompleted(event) => TurnCostObservationKind::ResponseCompleted { + response_id: event.response_id.clone(), + }, EventMsg::TurnComplete(_) => TurnCostObservationKind::Finished { interrupted: false }, EventMsg::TurnAborted(_) => TurnCostObservationKind::Finished { interrupted: true }, _ => return, @@ -193,9 +199,7 @@ impl TurnCostWorkerHandle { impl WorkerRuntime { async fn run(self, receiver: mpsc::Receiver, shutdown: CancellationToken) { let auth_changes = match &self.backend { - TurnCostBackend::OpenAiApiKey(auth_manager) => { - Some(auth_manager.auth_change_receiver()) - } + TurnCostBackend::OpenAi(auth_manager) => Some(auth_manager.auth_change_receiver()), TurnCostBackend::ModelProvider(_) => None, }; let backend_availability = self.probe_backend().await; @@ -225,6 +229,8 @@ impl WorkerRuntime { if changed.is_err() { break; } + // Cost telemetry is best effort: auth changes may drop pending turns. + // Callers can query the cost API directly to reconcile missing telemetry. self.turns.clear(); while receiver.try_recv().is_ok() {} backend_availability = self.probe_backend().await; @@ -260,26 +266,29 @@ impl WorkerRuntime { match tokio::time::timeout(REQUEST_TIMEOUT, self.query_turn_costs(&probe_turn_ids)).await { Ok(Ok(Some(_))) => BackendAvailability::Ready, Ok(Ok(None)) => match self.backend { - TurnCostBackend::OpenAiApiKey(_) => BackendAvailability::AwaitingAuthChange, + TurnCostBackend::OpenAi(_) => BackendAvailability::AwaitingAuthChange, TurnCostBackend::ModelProvider(_) => BackendAvailability::Disabled, }, Ok(Err(error)) => match error.status().map(|status| status.as_u16()) { - Some(401 | 403) if matches!(self.backend, TurnCostBackend::OpenAiApiKey(_)) => { + Some(401 | 403) if matches!(self.backend, TurnCostBackend::OpenAi(_)) => { tracing::debug!( - "turn cost worker waiting for auth change after backend availability check: {error}" + status = error.status().map(|status| status.as_u16()), + "turn cost worker waiting for auth change after backend availability check" ); BackendAvailability::AwaitingAuthChange } Some(401 | 403 | 429) => BackendAvailability::RetryProbe, Some(400..=499) => { tracing::debug!( - "turn cost worker disabled by backend availability check: {error}" + status = error.status().map(|status| status.as_u16()), + "turn cost worker disabled by backend availability check" ); BackendAvailability::Disabled } _ => { tracing::debug!( - "turn cost worker backend availability check failed temporarily: {error}" + status = error.status().map(|status| status.as_u16()), + "turn cost worker backend availability check failed temporarily" ); BackendAvailability::RetryProbe } @@ -302,18 +311,22 @@ impl WorkerRuntime { .or_insert(TurnCostEntry { thread_id: observation.thread_id, session_telemetry: *session_telemetry, - expected_response_count: 0, + expected_response_ids: HashSet::new(), status: TurnCostStatus::Running, next_poll_at: Instant::now(), attempt_count: 0, }); } } - TurnCostObservationKind::ResponseCompleted => { + TurnCostObservationKind::ResponseCompleted { response_id } => { if let Some(entry) = self.turns.get_mut(&observation.turn_id) && entry.status == TurnCostStatus::Running { - entry.expected_response_count = entry.expected_response_count.saturating_add(1); + if entry.expected_response_ids.len() >= MAX_TRACKED_RESPONSES_PER_TURN { + self.turns.remove(&observation.turn_id); + return; + } + entry.expected_response_ids.insert(response_id); } } TurnCostObservationKind::Finished { interrupted } => { @@ -345,22 +358,25 @@ impl WorkerRuntime { .map(|(turn_id, _)| turn_id.clone()) .collect(); if !due_turn_ids.is_empty() { - self.poll_api_key_entries(&due_turn_ids).await; + self.poll_entries(&due_turn_ids).await; } } - async fn poll_api_key_entries(&mut self, turn_ids: &[String]) { + async fn poll_entries(&mut self, turn_ids: &[String]) { let costs = match tokio::time::timeout(REQUEST_TIMEOUT, self.query_turn_costs(turn_ids)).await { Ok(Ok(Some(costs))) => costs, Ok(Ok(None)) => return, Ok(Err(error)) => { - warn!("failed to query API-key turn costs: {error}"); + warn!( + status = error.status().map(|status| status.as_u16()), + "failed to query turn costs" + ); self.retry_entries(turn_ids); return; } Err(_) => { - warn!("timed out querying API-key turn costs"); + warn!("timed out querying turn costs"); self.retry_entries(turn_ids); return; } @@ -374,7 +390,7 @@ impl WorkerRuntime { self.retry_entry(turn_id); continue; }; - self.process_api_key_cost(turn_id, cost); + self.process_turn_cost(turn_id, cost); } } @@ -383,10 +399,16 @@ impl WorkerRuntime { turn_ids: &[String], ) -> Result>, RequestError> { match &self.backend { - TurnCostBackend::OpenAiApiKey(auth_manager) => { + TurnCostBackend::OpenAi(auth_manager) => { let Some(auth) = auth_manager.auth().await else { return Ok(None); }; + if auth.is_chatgpt_auth() { + return self + .query_chatgpt_turn_costs(&auth, turn_ids) + .await + .map(Some); + } if !auth.is_api_key_auth() { return Ok(None); } @@ -436,7 +458,7 @@ impl WorkerRuntime { } } - fn process_api_key_cost(&mut self, turn_id: &str, cost: &ApiKeyTurnCost) { + fn process_turn_cost(&mut self, turn_id: &str, cost: &ApiKeyTurnCost) { if cost.status != ApiKeyTurnCostStatus::Priced { self.retry_entry(turn_id); return; @@ -454,7 +476,7 @@ impl WorkerRuntime { let Some(entry) = self.turns.get(turn_id) else { return; }; - if response_count < entry.expected_response_count { + if response_count < entry.expected_response_ids.len() as u64 { self.retry_entry(turn_id); return; } @@ -499,6 +521,9 @@ impl WorkerRuntime { } } +#[path = "turn_cost_worker_chatgpt.rs"] +mod chatgpt; + #[cfg(test)] #[path = "turn_cost_worker_tests.rs"] mod tests; diff --git a/codex-rs/app-server/src/turn_cost_worker_chatgpt.rs b/codex-rs/app-server/src/turn_cost_worker_chatgpt.rs new file mode 100644 index 0000000000..4c300e0bcc --- /dev/null +++ b/codex-rs/app-server/src/turn_cost_worker_chatgpt.rs @@ -0,0 +1,81 @@ +//! Adapts workspace-visible SiWC estimates to the shared turn-cost emission path. +//! Emit only nonnegative, visible amounts covering every observed completed response; +//! missing amounts/settlement data must never be interpreted as a zero-cost turn. + +use super::ApiKeyTurnCost; +use super::ApiKeyTurnCostStatus; +use super::BackendClient; +use super::RequestError; +use super::WorkerRuntime; +use codex_login::CodexAuth; +use std::collections::BTreeMap; + +impl WorkerRuntime { + pub(super) async fn query_chatgpt_turn_costs( + &self, + auth: &CodexAuth, + turn_ids: &[String], + ) -> Result, RequestError> { + let mut threads: BTreeMap> = BTreeMap::new(); + for turn_id in turn_ids { + // Availability probes use an untracked, random thread/turn pair. + let thread_id = self + .turns + .get(turn_id) + .map(|entry| entry.thread_id) + .unwrap_or_default(); + threads + .entry(thread_id.to_string()) + .or_default() + .push(turn_id.clone()); + } + let client = BackendClient::from_auth( + self.config.chatgpt_base_url.clone(), + auth, + self.config.http_client_factory(), + ); + let costs = client.query_chatgpt_turn_costs(&threads).await?; + let mut settled_costs = Vec::new(); + for thread in costs { + let Some(requested_turns) = threads.get(&thread.thread_id) else { + continue; + }; + for cost in thread.turns { + if !requested_turns.contains(&cost.turn_id) { + continue; + } + let Some(entry) = self.turns.get(&cost.turn_id) else { + continue; + }; + let (Some(micros), Some(settled_ids)) = + (cost.estimated_usage_usd_micros, cost.settled_response_ids) + else { + continue; + }; + // Missing/hidden dollars are not zero. A projection can appear before + // all responses settle, so require every locally observed response. + if micros < 0 + || entry.expected_response_ids.is_empty() + || !entry + .expected_response_ids + .iter() + .all(|id| settled_ids.contains(id)) + { + continue; + } + // Reuse the existing emission path without a floating-point conversion. + settled_costs.push(ApiKeyTurnCost { + turn_id: cost.turn_id, + status: ApiKeyTurnCostStatus::Priced, + total_usd: Some(format!("{}.{:06}", micros / 1_000_000, micros % 1_000_000)), + event_count: Some(settled_ids.len() as u64), + responses: None, + model: cost.model, + speed: None, + reasoning_effort: None, + }); + } + } + Ok(settled_costs) + } +} diff --git a/codex-rs/app-server/src/turn_cost_worker_chatgpt_tests.rs b/codex-rs/app-server/src/turn_cost_worker_chatgpt_tests.rs new file mode 100644 index 0000000000..432641bef8 --- /dev/null +++ b/codex-rs/app-server/src/turn_cost_worker_chatgpt_tests.rs @@ -0,0 +1,73 @@ +//! Checks unusable SiWC estimates and exact conversion; public app-server/OTLP coverage lives in turn_cost_otel. + +use super::*; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use test_case::test_case; + +#[test_case(json!(null), json!(["resp-one", "resp-two"]), None; "hidden or incomplete amount")] +#[test_case(json!(-1), json!(["resp-one", "resp-two"]), None; "negative amount")] +#[test_case(json!(1), json!(null), None; "missing settlement")] +#[test_case(json!(1), json!([]), None; "unsettled")] +#[test_case(json!(1), json!(["unrelated"]), None; "wrong response")] +#[test_case(json!(1), json!(["resp-one"]), None; "only one response settled")] +#[test_case(json!(0), json!(["resp-one", "resp-two"]), Some("0.000000"); "zero")] +#[test_case(json!(i64::MAX), json!(["resp-one", "resp-two"]), Some("9223372036854.775807"); "integer precision")] +#[tokio::test] +async fn chatgpt_cost_requires_visible_amount_and_matching_settlement( + micros: Value, + settled_ids: Value, + expected: Option<&str>, +) { + let server = MockServer::start().await; + let mut runtime = test_runtime( + &server, + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()), + ) + .await; + let thread_id = ThreadId::new(); + runtime.turns.insert( + "turn-1".to_string(), + TurnCostEntry { + thread_id, + session_telemetry: test_session_telemetry(thread_id), + expected_response_ids: HashSet::from(["resp-one".to_string(), "resp-two".to_string()]), + status: TurnCostStatus::Completed, + next_poll_at: Instant::now(), + attempt_count: 0, + }, + ); + Mock::given(method("POST")) + .and(path("/api/codex/usage/thread-estimates/query")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"threads": [{ + "thread_id": thread_id, "turns": [{"turn_id": "turn-1", + "estimated_usage_usd_micros": micros, "settled_response_ids": settled_ids}] + }]})), + ) + .expect(1) + .mount(&server) + .await; + let costs = runtime + .query_turn_costs(&["turn-1".to_string()]) + .await + .expect("query") + .expect("SiWC enabled"); + assert_eq!( + costs, + expected + .map(|amount| ApiKeyTurnCost { + turn_id: "turn-1".to_string(), + status: ApiKeyTurnCostStatus::Priced, + total_usd: Some(amount.to_string()), + event_count: Some(2), + responses: None, + model: None, + speed: None, + reasoning_effort: None, + }) + .into_iter() + .collect::>() + ); +} diff --git a/codex-rs/app-server/src/turn_cost_worker_tests.rs b/codex-rs/app-server/src/turn_cost_worker_tests.rs index d7b8b2a3d2..c1b62699f7 100644 --- a/codex-rs/app-server/src/turn_cost_worker_tests.rs +++ b/codex-rs/app-server/src/turn_cost_worker_tests.rs @@ -11,7 +11,9 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::TurnStartedEvent; use pretty_assertions::assert_eq; use tempfile::TempDir; +use test_case::test_case; use tokio::time::timeout; +use tracing::instrument::WithSubscriber; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; @@ -21,6 +23,9 @@ use wiremock::matchers::path; const TURN_COST_PATH: &str = "/v1/analytics/codex/turn-costs"; +#[path = "turn_cost_worker_chatgpt_tests.rs"] +mod chatgpt; + #[tokio::test] async fn worker_starts_with_otlp_metrics_exporter_without_log_exporter() { let server = MockServer::start().await; @@ -308,7 +313,9 @@ async fn priced_cost_uses_telemetry_captured_before_thread_removal() { runtime.record_observation(TurnCostObservation { thread_id, turn_id: turn_id.to_string(), - kind: TurnCostObservationKind::ResponseCompleted, + kind: TurnCostObservationKind::ResponseCompleted { + response_id: "resp-one".to_string(), + }, }); runtime.record_observation(TurnCostObservation { thread_id, @@ -316,7 +323,7 @@ async fn priced_cost_uses_telemetry_captured_before_thread_removal() { kind: TurnCostObservationKind::Finished { interrupted: false }, }); - runtime.process_api_key_cost( + runtime.process_turn_cost( turn_id, &ApiKeyTurnCost { turn_id: turn_id.to_string(), @@ -346,7 +353,7 @@ async fn priced_cost_waits_for_every_response_when_response_costs_are_available( TurnCostEntry { thread_id, session_telemetry: test_session_telemetry(thread_id), - expected_response_count: 2, + expected_response_ids: HashSet::from(["resp-one".to_string(), "resp-two".to_string()]), status: TurnCostStatus::Completed, next_poll_at: Instant::now(), attempt_count: 0, @@ -366,7 +373,7 @@ async fn priced_cost_waits_for_every_response_when_response_costs_are_available( speed: Some("fast".to_string()), reasoning_effort: Some("high".to_string()), }; - runtime.process_api_key_cost(turn_id, &cost); + runtime.process_turn_cost(turn_id, &cost); let entry = runtime.turns.get(turn_id).expect("turn remains tracked"); assert_eq!(entry.attempt_count, 1); @@ -379,7 +386,7 @@ async fn priced_cost_waits_for_every_response_when_response_costs_are_available( response_id: "resp-two".to_string(), total_usd: "0.50".to_string(), }); - runtime.process_api_key_cost(turn_id, &cost); + runtime.process_turn_cost(turn_id, &cost); assert_eq!(runtime.turns.len(), 0); } @@ -392,7 +399,7 @@ async fn test_runtime(server: &MockServer, auth_manager: Arc) -> Wo .await .expect("test config"); config.chatgpt_base_url = server.uri(); - let backend = TurnCostBackend::OpenAiApiKey(Arc::clone(&auth_manager)); + let backend = TurnCostBackend::OpenAi(Arc::clone(&auth_manager)); WorkerRuntime { config: Arc::new(config), backend, @@ -443,3 +450,43 @@ async fn wait_for_request_count(server: &MockServer, expected: usize) { .await .expect("timed out waiting for turn-cost request"); } + +#[test_case(500; "HTTP failure")] +#[test_case(200; "decode failure")] +#[tokio::test] +async fn turn_cost_failures_do_not_log_response_bodies(status: u16) { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(TURN_COST_PATH)) + .respond_with( + ResponseTemplate::new(status).set_body_string("private-workspace-response-marker"), + ) + .expect(2) + .mount(&server) + .await; + let mut runtime = test_runtime( + &server, + AuthManager::from_auth_for_testing(CodexAuth::from_api_key("sk-test")), + ) + .await; + let capture = tempfile::NamedTempFile::new().expect("log capture"); + let subscriber = tracing_subscriber::fmt() + .with_ansi(false) + .with_max_level(tracing::Level::DEBUG) + .with_writer(Arc::new(capture.reopen().expect("log writer"))) + .finish(); + async { + assert_eq!( + runtime.probe_backend().await, + BackendAvailability::RetryProbe + ); + runtime.poll_entries(&["turn-1".to_string()]).await; + } + .with_subscriber(subscriber) + .await; + let logs = std::fs::read_to_string(capture.path()).expect("captured logs"); + assert!(logs.contains("backend availability check failed temporarily")); + assert!(logs.contains("failed to query turn costs")); + assert!(!logs.contains("private-workspace-response-marker")); + server.verify().await; +} diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 7ea4813a71..23c2938eba 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -115,6 +115,7 @@ mod thread_status; mod thread_timeline; mod thread_unarchive; mod thread_unsubscribe; +mod turn_cost_otel; mod turn_interrupt; mod turn_settings_update; mod turn_start; diff --git a/codex-rs/app-server/tests/suite/v2/turn_cost_otel.rs b/codex-rs/app-server/tests/suite/v2/turn_cost_otel.rs new file mode 100644 index 0000000000..80e3f2eabf --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/turn_cost_otel.rs @@ -0,0 +1,323 @@ +//! Exercises the best-effort SiWC cost happy path through real app-server startup and OTLP. +//! The public in-process transport permits advancing the production timer without test-only hooks. + +use anyhow::Result; +use app_test_support::ChatGptAuthFixture; +use app_test_support::write_chatgpt_auth; +use app_test_support::write_models_cache; +use codex_app_server::in_process; +use codex_app_server::in_process::InProcessServerEvent; +use codex_app_server::in_process::InProcessStartArgs; +use codex_app_server_protocol::ClientInfo; +use codex_app_server_protocol::ClientRequest; +use codex_app_server_protocol::InitializeCapabilities; +use codex_app_server_protocol::InitializeParams; +use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ServerNotification; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::ThreadStartResponse; +use codex_app_server_protocol::TurnEnvironmentParams; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; +use codex_app_server_protocol::UserInput; +use codex_arg0::Arg0DispatchPaths; +use codex_config::CloudConfigBundleLoader; +use codex_config::LoaderOverrides; +use codex_config::NoopThreadConfigLoader; +use codex_core::config::ConfigBuilder; +use codex_exec_server::EnvironmentManager; +use codex_feedback::CodexFeedback; +use codex_login::AuthCredentialsStoreMode; +use codex_protocol::protocol::SessionSource; +use core_test_support::responses; +use core_test_support::test_codex::test_env; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use std::sync::Arc; +use std::time::Duration; +use tempfile::TempDir; +use tokio::time::timeout; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +const COST_PATH: &str = "/api/codex/usage/thread-estimates/query"; + +#[tokio::test] +async fn chatgpt_turn_cost_reaches_otlp_on_success() -> Result<()> { + let server = MockServer::start().await; + let collector = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/metrics")) + .respond_with(ResponseTemplate::new(200)) + .mount(&collector) + .await; + Mock::given(method("POST")).and(path(COST_PATH)) + .and(header("authorization", "Bearer siwc-token")) + .and(header("chatgpt-account-id", "workspace-a")) + .respond_with(move |request: &wiremock::Request| { + let body: Value = serde_json::from_slice(&request.body).expect("cost request"); + assert_eq!(body["include_settled_response_ids"], true); + let threads: Vec<_> = body["threads"].as_array().expect("threads").iter().map(|thread| { + let turns: Vec<_> = thread["turn_ids"].as_array().expect("turns").iter().map(|turn_id| json!({ + "turn_id": turn_id, "model": "mock-model", "estimated_usage_usd_micros": 1250001, + "settled_response_ids": ["resp-cost"] + })).collect(); + json!({"thread_id": thread["thread_id"], "turns": turns}) + }).collect(); + ResponseTemplate::new(200).set_body_json(json!({"threads": threads})) + }).mount(&server).await; + let model = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-cost"), + responses::ev_assistant_message("msg-cost", "Done"), + responses::ev_completed("resp-cost"), + ]), + ) + .await; + let home = TempDir::new()?; + std::fs::write( + home.path().join("config.toml"), + format!( + r#" +model = "mock-model" +model_provider = "openai" +openai_base_url = "{}/v1" +chatgpt_base_url = "{}" +approval_policy = "never" +[features] +responses_websockets = false +responses_websockets_v2 = false +runtime_metrics = true +[analytics] +enabled = true +[otel] +exporter = "none" +metrics_exporter = {{ otlp-http = {{ endpoint = "{}/metrics", protocol = "json" }} }} +"#, + server.uri(), + server.uri(), + collector.uri() + ), + )?; + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("siwc-token") + .account_id("workspace-a") + .chatgpt_account_id("workspace-a") + .chatgpt_user_id("user-a"), + AuthCredentialsStoreMode::File, + )?; + write_models_cache(home.path())?; + let loader_overrides = LoaderOverrides::without_managed_config_for_tests(); + let config = Arc::new( + ConfigBuilder::default() + .codex_home(home.path().to_path_buf()) + .loader_overrides(loader_overrides.clone()) + .build() + .await?, + ); + let provider = codex_core::otel_init::build_provider( + &config, + "test", + Some("codex-app-server"), + /*default_analytics_enabled*/ false, + ) + .map_err(|error| anyhow::anyhow!(error.to_string()))? + .expect("OTLP enabled"); + let metrics = provider.metrics.clone().expect("metrics exporter"); + let env = test_env().await?; + let environment_manager = if let Some(url) = env.exec_server_url() { + EnvironmentManager::create_for_tests( + Some(url.to_string()), + /*local_runtime_paths*/ None, + ) + .await + } else { + EnvironmentManager::default_for_tests() + }; + let mut client = in_process::start(InProcessStartArgs { + arg0_paths: Arg0DispatchPaths::default(), + config, + cli_overrides: Vec::new(), + loader_overrides, + strict_config: false, + cloud_config_bundle: CloudConfigBundleLoader::default(), + thread_config_loader: Arc::new(NoopThreadConfigLoader), + feedback: CodexFeedback::new(), + log_db: None, + state_db: None, + environment_manager: Arc::new(environment_manager), + config_warnings: Vec::new(), + session_source: SessionSource::Cli, + enable_codex_api_key_env: false, + initialize: InitializeParams { + client_info: ClientInfo { + name: "codex-app-server-tests".to_string(), + title: None, + version: "test".to_string(), + }, + capabilities: Some(InitializeCapabilities { + experimental_api: true, + ..Default::default() + }), + }, + channel_capacity: in_process::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, + }) + .await?; + wait_for_cost_queries(&server, /*count*/ 1).await; + let response = client + .request(ClientRequest::ThreadStart { + request_id: RequestId::Integer(1), + params: ThreadStartParams { + environments: Some(vec![TurnEnvironmentParams { + environment_id: env.selection().environment_id.clone(), + cwd: env.selection().cwd.clone().into(), + runtime_workspace_roots: None, + }]), + ..Default::default() + }, + }) + .await? + .expect("thread/start"); + let ThreadStartResponse { thread, .. } = serde_json::from_value(response)?; + let response = client + .request(ClientRequest::TurnStart { + request_id: RequestId::Integer(2), + params: TurnStartParams { + thread_id: thread.id.clone(), + input: vec![UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, + }) + .await? + .expect("turn/start"); + let TurnStartResponse { turn } = serde_json::from_value(response)?; + timeout(Duration::from_secs(/*secs*/ 30), async { + while let Some(event) = client.next_event().await { + if let InProcessServerEvent::ServerNotification(event) = event + && let ServerNotification::TurnCompleted(event) = *event + { + assert!( + event.turn.error.is_none(), + "turn failed: {:?}", + event.turn.error + ); + break; + } + } + }) + .await?; + assert_eq!( + model.single_request().header("authorization").as_deref(), + Some("Bearer siwc-token") + ); + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(/*secs*/ 301)).await; + tokio::time::resume(); + wait_for_cost_queries(&server, /*count*/ 2).await; + // Synchronize with actual recording, not just the backend receiving a request. + timeout(Duration::from_secs(/*secs*/ 10), async { + loop { + if metrics + .snapshot() + .expect("snapshot") + .scope_metrics() + .flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics) + .any(|metric| metric.name() == "codex.turn.cost_microusd") + { + break; + } + tokio::task::yield_now().await; + } + }) + .await?; + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(/*secs*/ 301)).await; + tokio::time::resume(); + client.shutdown().await?; + tokio::task::spawn_blocking(move || provider.shutdown()).await?; + let mut points = Vec::new(); + for request in collector.received_requests().await.expect("OTLP requests") { + let body: Value = serde_json::from_slice(&request.body)?; + for resource in body["resourceMetrics"] + .as_array() + .expect("resource metrics") + { + for scope in resource["scopeMetrics"].as_array().expect("scope metrics") { + for metric in scope["metrics"].as_array().expect("metrics") { + if metric["name"] == "codex.turn.cost_microusd" { + points.extend( + metric["sum"]["dataPoints"] + .as_array() + .expect("cost points") + .clone(), + ); + } + } + } + } + } + assert_eq!(points.len(), 1); + assert_eq!(points[0]["asInt"], json!(1250001)); + let attrs: std::collections::BTreeMap<_, _> = points[0]["attributes"] + .as_array() + .expect("attributes") + .iter() + .map(|attr| { + ( + attr["key"].as_str().expect("key"), + attr["value"]["stringValue"].as_str().expect("value"), + ) + }) + .collect(); + assert_eq!(attrs["turn.id"], turn.id); + assert_eq!(attrs["conversation.id"], thread.id); + assert_eq!(attrs["auth_mode"], "Chatgpt"); + let requests = server + .received_requests() + .await + .expect("settlement requests"); + let cost_requests: Vec<_> = requests + .iter() + .filter(|request| request.url.path() == COST_PATH) + .collect(); + assert_eq!(cost_requests.len(), 2); + assert_eq!( + serde_json::from_slice::(&cost_requests[1].body)?, + json!({ + "threads": [{"thread_id": thread.id, "turn_ids": [turn.id]}], "include_settled_response_ids": true + }) + ); + server.verify().await; + Ok(()) +} + +async fn wait_for_cost_queries(server: &MockServer, count: usize) { + timeout(Duration::from_secs(/*secs*/ 15), async { + loop { + if server + .received_requests() + .await + .expect("requests") + .iter() + .filter(|request| request.url.path() == COST_PATH) + .count() + >= count + { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("cost query"); +} diff --git a/codex-rs/backend-client/src/client.rs b/codex-rs/backend-client/src/client.rs index 2da9be27b4..0c4f96395f 100644 --- a/codex-rs/backend-client/src/client.rs +++ b/codex-rs/backend-client/src/client.rs @@ -34,10 +34,13 @@ use serde::Serialize; use serde::de::DeserializeOwned; use std::fmt; +mod chatgpt_turn_cost; mod rate_limit_resets; mod thread_usage; pub(crate) mod turn_usage; +pub use chatgpt_turn_cost::ChatgptThreadTurnCosts; +pub use chatgpt_turn_cost::ChatgptTurnCost; pub use thread_usage::ThreadUsage; pub use thread_usage::ThreadUsageBreakdownGroup; diff --git a/codex-rs/backend-client/src/client/chatgpt_turn_cost.rs b/codex-rs/backend-client/src/client/chatgpt_turn_cost.rs new file mode 100644 index 0000000000..a0398ef5c0 --- /dev/null +++ b/codex-rs/backend-client/src/client/chatgpt_turn_cost.rs @@ -0,0 +1,81 @@ +//! Queries SiWC per-turn estimates using the client's Codex or WHAM route and auth. +//! Always request settled response IDs and retain thread/turn association and missing amounts. + +use super::Client; +use super::PathStyle; +use super::RequestError; +use http::Method; +use serde::Deserialize; +use serde::Serialize; +use std::collections::BTreeMap; + +/// Workspace-visible estimated cost and settled responses for a ChatGPT turn. +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +pub struct ChatgptTurnCost { + pub turn_id: String, + pub model: Option, + pub estimated_usage_usd_micros: Option, + pub settled_response_ids: Option>, +} + +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +pub struct ChatgptThreadTurnCosts { + pub thread_id: String, + pub turns: Vec, +} + +#[derive(Serialize)] +struct ThreadTurnIds<'a> { + thread_id: &'a str, + turn_ids: &'a [String], +} + +#[derive(Serialize)] +struct TurnCostsRequest<'a> { + threads: Vec>, + include_settled_response_ids: bool, +} + +#[derive(Deserialize)] +struct TurnCostsResponse { + threads: Vec, +} + +impl Client { + /// Queries per-turn estimates using the signed-in user's workspace permissions. + pub async fn query_chatgpt_turn_costs( + &self, + threads: &BTreeMap>, + ) -> Result, RequestError> { + let url = match self.path_style { + PathStyle::CodexApi => { + format!("{}/api/codex/usage/thread-estimates/query", self.base_url) + } + PathStyle::ChatGptApi => { + format!("{}/wham/usage/thread-estimates/query", self.base_url) + } + }; + let request = self + .request(Method::POST, &url) + .headers(self.headers()) + .json(&TurnCostsRequest { + threads: threads + .iter() + .map(|(thread_id, turn_ids)| ThreadTurnIds { + thread_id, + turn_ids, + }) + .collect(), + include_settled_response_ids: true, + }); + let (body, content_type) = self.exec_request_detailed(request, "POST", &url).await?; + let response: TurnCostsResponse = self + .decode_json(&url, &content_type, &body) + .map_err(RequestError::Other)?; + Ok(response.threads) + } +} + +#[cfg(test)] +#[path = "chatgpt_turn_cost_tests.rs"] +mod tests; diff --git a/codex-rs/backend-client/src/client/chatgpt_turn_cost_tests.rs b/codex-rs/backend-client/src/client/chatgpt_turn_cost_tests.rs new file mode 100644 index 0000000000..e6c8fadb35 --- /dev/null +++ b/codex-rs/backend-client/src/client/chatgpt_turn_cost_tests.rs @@ -0,0 +1,108 @@ +//! Exercises both authenticated SiWC routes and preserves nullable/omitted estimate fields. + +use super::*; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; +use codex_login::CodexAuth; +use pretty_assertions::assert_eq; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_json; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +#[tokio::test] +async fn chatgpt_turn_costs_use_workspace_auth_and_chatgpt_path() { + check_query( + "/backend-api", + "/backend-api/wham/usage/thread-estimates/query", + ) + .await; +} + +#[tokio::test] +async fn chatgpt_turn_costs_use_workspace_auth_and_codex_path() { + check_query("", "/api/codex/usage/thread-estimates/query").await; +} + +async fn check_query(base_path: &str, endpoint: &str) { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(endpoint)) + .and(header("authorization", "Bearer Access Token")) + .and(header("chatgpt-account-id", "account_id")) + .and(header("content-type", "application/json")) + .and(body_json(serde_json::json!({ + "threads": [ + {"thread_id": "thread-a", "turn_ids": ["turn-a"]}, + {"thread_id": "thread-b", "turn_ids": ["turn-b", "turn-c"]} + ], + "include_settled_response_ids": true + }))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "threads": [ + {"thread_id": "thread-a", "turns": [{ + "turn_id": "turn-a", "model": "gpt-5.6", + "estimated_usage_usd_micros": 1250001, + "settled_response_ids": ["resp-a", "resp-b"] + }]}, + {"thread_id": "thread-b", "turns": [ + {"turn_id": "turn-b", "estimated_usage_usd_micros": null, + "settled_response_ids": null}, + {"turn_id": "turn-c", "estimated_usage_usd_micros": 0} + ]} + ] + }))) + .expect(1) + .mount(&server) + .await; + let client = Client::from_auth( + format!("{}{base_path}", server.uri()), + &CodexAuth::create_dummy_chatgpt_auth_for_testing(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); + let result = client + .query_chatgpt_turn_costs(&BTreeMap::from([ + ("thread-a".to_string(), vec!["turn-a".to_string()]), + ( + "thread-b".to_string(), + vec!["turn-b".to_string(), "turn-c".to_string()], + ), + ])) + .await + .expect("ChatGPT turn costs"); + assert_eq!( + result, + vec![ + ChatgptThreadTurnCosts { + thread_id: "thread-a".to_string(), + turns: vec![ChatgptTurnCost { + turn_id: "turn-a".to_string(), + model: Some("gpt-5.6".to_string()), + estimated_usage_usd_micros: Some(1250001), + settled_response_ids: Some(vec!["resp-a".to_string(), "resp-b".to_string()]), + }], + }, + ChatgptThreadTurnCosts { + thread_id: "thread-b".to_string(), + turns: vec![ + ChatgptTurnCost { + turn_id: "turn-b".to_string(), + model: None, + estimated_usage_usd_micros: None, + settled_response_ids: None, + }, + ChatgptTurnCost { + turn_id: "turn-c".to_string(), + model: None, + estimated_usage_usd_micros: Some(0), + settled_response_ids: None, + }, + ], + }, + ] + ); + server.verify().await; +} diff --git a/codex-rs/backend-client/src/lib.rs b/codex-rs/backend-client/src/lib.rs index 1bc92cc49c..4329da770b 100644 --- a/codex-rs/backend-client/src/lib.rs +++ b/codex-rs/backend-client/src/lib.rs @@ -2,6 +2,8 @@ mod client; pub(crate) mod types; pub use client::AddCreditsNudgeCreditType; +pub use client::ChatgptThreadTurnCosts; +pub use client::ChatgptTurnCost; pub use client::Client; pub use client::RequestError; pub use client::ThreadUsage;