diff --git a/codex-rs/backend-client/src/client.rs b/codex-rs/backend-client/src/client.rs index da9b0a3d3f..6de0df4211 100644 --- a/codex-rs/backend-client/src/client.rs +++ b/codex-rs/backend-client/src/client.rs @@ -35,6 +35,10 @@ use serde::de::DeserializeOwned; use std::fmt; mod rate_limit_resets; +mod thread_usage; + +pub use thread_usage::ThreadUsage; +pub use thread_usage::ThreadUsageBreakdownGroup; #[derive(Debug)] pub enum RequestError { diff --git a/codex-rs/backend-client/src/client/thread_usage.rs b/codex-rs/backend-client/src/client/thread_usage.rs new file mode 100644 index 0000000000..9e879f7f1d --- /dev/null +++ b/codex-rs/backend-client/src/client/thread_usage.rs @@ -0,0 +1,88 @@ +//! Authoritative estimated credit and dollar usage for an individual Codex thread. + +use super::Client; +use super::PathStyle; +use super::RequestError; +use anyhow::anyhow; +use http::Method; +use http::header::CONTENT_TYPE; +use http::header::HeaderValue; +use serde::Deserialize; +use serde::Serialize; + +/// Backend usage grouped by model, reasoning effort, and response speed. +#[derive(Clone, Debug, Deserialize, PartialEq, Eq)] +pub struct ThreadUsageBreakdownGroup { + pub model: Option, + pub reasoning_effort: Option, + pub speed: Option, + pub estimated_usage_credits_micros: i64, + pub net_new_input_tokens: Option, + pub cached_input_tokens: Option, + pub input_tokens: Option, + pub output_tokens: Option, + pub total_tokens: Option, +} + +/// Backend-estimated usage totals expressed in integer millionths. +#[derive(Clone, Debug, Deserialize, PartialEq, Eq)] +pub struct ThreadUsage { + pub thread_id: String, + pub estimated_usage_credits_micros: i64, + pub estimated_usage_usd_micros: Option, + #[serde(default)] + pub groups: Vec, +} + +#[derive(Serialize)] +struct ThreadUsageQueryRequest<'a> { + thread_ids: [&'a str; 1], +} + +#[derive(Deserialize)] +struct ThreadUsageQueryResponse { + threads: Vec, +} + +impl Client { + /// Reads authoritative estimated totals without maintaining a second usage ledger. + pub async fn get_thread_usage(&self, thread_id: &str) -> Result { + let url = self.thread_usage_url(); + let request = self + .request(Method::POST, &url) + .headers(self.headers()) + .header(CONTENT_TYPE, HeaderValue::from_static("application/json")) + .json(&ThreadUsageQueryRequest { + thread_ids: [thread_id], + }); + let (body, content_type) = self.exec_request_detailed(request, "POST", &url).await?; + let response = self + .decode_json::(&url, &content_type, &body) + .map_err(RequestError::from)?; + + response + .threads + .into_iter() + .find(|usage| usage.thread_id == thread_id) + .ok_or_else(|| { + RequestError::from(anyhow!( + "thread usage response did not contain requested thread {thread_id}" + )) + }) + } + + fn thread_usage_url(&self) -> String { + match self.path_style { + PathStyle::CodexApi => { + format!("{}/api/codex/usage/thread_usage/query", self.base_url) + } + PathStyle::ChatGptApi => { + format!("{}/wham/usage/thread_usage/query", self.base_url) + } + } + } +} + +#[cfg(test)] +#[path = "thread_usage_tests.rs"] +mod tests; diff --git a/codex-rs/backend-client/src/client/thread_usage_tests.rs b/codex-rs/backend-client/src/client/thread_usage_tests.rs new file mode 100644 index 0000000000..d8dc620856 --- /dev/null +++ b/codex-rs/backend-client/src/client/thread_usage_tests.rs @@ -0,0 +1,146 @@ +use super::*; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; +use pretty_assertions::assert_eq; +use serde_json::json; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_json; +use wiremock::matchers::method; +use wiremock::matchers::path; + +#[test] +fn thread_usage_contract_uses_expected_paths_and_payload() { + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + assert_eq!( + Client::new("https://example.test", factory.clone()).thread_usage_url(), + "https://example.test/api/codex/usage/thread_usage/query" + ); + assert_eq!( + Client::new("https://chatgpt.com/backend-api", factory).thread_usage_url(), + "https://chatgpt.com/backend-api/wham/usage/thread_usage/query" + ); + assert_eq!( + serde_json::to_value(ThreadUsageQueryRequest { + thread_ids: ["thread-123"], + }) + .expect("serialize thread usage request"), + json!({ "thread_ids": ["thread-123"] }) + ); +} + +#[tokio::test] +async fn get_thread_usage_returns_requested_thread_totals() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/codex/usage/thread_usage/query")) + .and(body_json(json!({ "thread_ids": ["thread-123"] }))) + .respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + "threads": [{ + "thread_id": "thread-123", + "estimated_usage_credits_micros": 46_000_000, + "estimated_usage_usd_micros": 1_820_000, + "groups": [{ + "model": "gpt-5.4", + "reasoning_effort": "high", + "speed": "fast", + "estimated_usage_credits_micros": 46_000_000, + "net_new_input_tokens": 80, + "cached_input_tokens": 20, + "input_tokens": 100, + "output_tokens": 40, + "total_tokens": 140 + }] + }] + }))) + .expect(/*r*/ 1) + .mount(&server) + .await; + + let client = Client::new( + server.uri(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); + assert_eq!( + client + .get_thread_usage("thread-123") + .await + .expect("read thread usage"), + ThreadUsage { + thread_id: "thread-123".to_string(), + estimated_usage_credits_micros: 46_000_000, + estimated_usage_usd_micros: Some(1_820_000), + groups: vec![ThreadUsageBreakdownGroup { + model: Some("gpt-5.4".to_string()), + reasoning_effort: Some("high".to_string()), + speed: Some("fast".to_string()), + estimated_usage_credits_micros: 46_000_000, + net_new_input_tokens: Some(80), + cached_input_tokens: Some(20), + input_tokens: Some(100), + output_tokens: Some(40), + total_tokens: Some(140), + }], + } + ); +} + +#[tokio::test] +async fn get_thread_usage_accepts_credits_without_usd_estimate() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/codex/usage/thread_usage/query")) + .respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + "threads": [{ + "thread_id": "thread-123", + "estimated_usage_credits_micros": 46_000_000, + "estimated_usage_usd_micros": null + }] + }))) + .expect(/*r*/ 1) + .mount(&server) + .await; + + let client = Client::new( + server.uri(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); + assert_eq!( + client + .get_thread_usage("thread-123") + .await + .expect("read credits without a dollar estimate"), + ThreadUsage { + thread_id: "thread-123".to_string(), + estimated_usage_credits_micros: 46_000_000, + estimated_usage_usd_micros: None, + groups: Vec::new(), + } + ); +} + +#[tokio::test] +async fn get_thread_usage_rejects_totals_for_another_thread() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + "threads": [{ + "thread_id": "another-thread", + "estimated_usage_credits_micros": 1, + "estimated_usage_usd_micros": 1 + }] + }))) + .mount(&server) + .await; + + let client = Client::new( + server.uri(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); + let error = client + .get_thread_usage("thread-123") + .await + .expect_err("reject usage for a different thread"); + assert!(error.to_string().contains("requested thread thread-123")); +} diff --git a/codex-rs/backend-client/src/lib.rs b/codex-rs/backend-client/src/lib.rs index 91aa10488d..9963687b0f 100644 --- a/codex-rs/backend-client/src/lib.rs +++ b/codex-rs/backend-client/src/lib.rs @@ -4,6 +4,8 @@ pub(crate) mod types; pub use client::AddCreditsNudgeCreditType; pub use client::Client; pub use client::RequestError; +pub use client::ThreadUsage; +pub use client::ThreadUsageBreakdownGroup; pub use types::AccountEntry; pub use types::AccountsCheckResponse; pub use types::CodeTaskDetailsResponse;