Emit turn cost telemetry for ChatGPT sessions (#41944)

## What changed

- Query workspace-visible turn estimates for ChatGPT-authenticated sessions through the appropriate Codex or ChatGPT endpoint.
- Track completed response IDs and emit `codex.turn.cost_microusd` only when the estimate is nonnegative, visible, and covers every observed response.
- Keep error response bodies out of turn-cost worker logs.

## Testing

- Cover both endpoint styles, nullable estimates, settlement validation, exact micro-USD conversion, and response-ID tracking.
- Add an app-server integration test that verifies a settled ChatGPT turn reaches the OTLP metrics exporter with its turn, conversation, and auth-mode attributes.

GitOrigin-RevId: 6ae3f8df729fad52984a0c6efafd8c96c7214e55
This commit is contained in:
xli-oai
2026-09-01 00:30:55 +00:00
committed by copyberry
parent 9c7edd4bc3
commit e39ab0c185
10 changed files with 776 additions and 32 deletions
+51 -26
View File
@@ -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<SessionTelemetry>,
},
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<String>,
status: TurnCostStatus,
next_poll_at: Instant,
attempt_count: u8,
@@ -81,7 +85,7 @@ struct WorkerRuntime {
#[derive(Clone)]
enum TurnCostBackend {
OpenAiApiKey(Arc<AuthManager>),
OpenAi(Arc<AuthManager>),
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<TurnCostObservation>, 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<Option<Vec<ApiKeyTurnCost>>, 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;
@@ -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<Vec<ApiKeyTurnCost>, RequestError> {
let mut threads: BTreeMap<String, Vec<String>> = 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)
}
}
@@ -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::<Vec<_>>()
);
}
@@ -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<AuthManager>) -> 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;
}
@@ -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;
@@ -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::<Value>(&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");
}
+3
View File
@@ -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;
@@ -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<String>,
pub estimated_usage_usd_micros: Option<i64>,
pub settled_response_ids: Option<Vec<String>>,
}
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
pub struct ChatgptThreadTurnCosts {
pub thread_id: String,
pub turns: Vec<ChatgptTurnCost>,
}
#[derive(Serialize)]
struct ThreadTurnIds<'a> {
thread_id: &'a str,
turn_ids: &'a [String],
}
#[derive(Serialize)]
struct TurnCostsRequest<'a> {
threads: Vec<ThreadTurnIds<'a>>,
include_settled_response_ids: bool,
}
#[derive(Deserialize)]
struct TurnCostsResponse {
threads: Vec<ChatgptThreadTurnCosts>,
}
impl Client {
/// Queries per-turn estimates using the signed-in user's workspace permissions.
pub async fn query_chatgpt_turn_costs(
&self,
threads: &BTreeMap<String, Vec<String>>,
) -> Result<Vec<ChatgptThreadTurnCosts>, 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;
@@ -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;
}
+2
View File
@@ -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;