File size: 4,459 Bytes
afa0cbf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | //! Authoritative estimated credit and dollar usage for bounded batches of Codex threads.
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<String>,
pub reasoning_effort: Option<String>,
pub speed: Option<String>,
pub estimated_usage_credits_micros: i64,
pub net_new_input_tokens: Option<i64>,
pub cached_input_tokens: Option<i64>,
pub input_tokens: Option<i64>,
pub output_tokens: Option<i64>,
pub total_tokens: Option<i64>,
}
/// 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<i64>,
#[serde(default)]
pub groups: Vec<ThreadUsageBreakdownGroup>,
}
#[derive(Serialize)]
struct ThreadUsageQueryRequest<'a> {
thread_ids: &'a [&'a str],
}
#[derive(Deserialize)]
struct ThreadUsageQueryResponse {
threads: Vec<ThreadUsageEstimate>,
}
// Preserve unavailable rows until all returned IDs have been validated.
#[derive(Deserialize)]
struct ThreadUsageEstimate {
thread_id: String,
estimated_usage_credits_micros: Option<i64>,
estimated_usage_usd_micros: Option<i64>,
groups: Option<Vec<ThreadUsageBreakdownGroup>>,
}
impl Client {
/// Reads authoritative estimated totals without maintaining a second usage ledger.
pub async fn get_thread_usage(&self, thread_id: &str) -> Result<ThreadUsage, RequestError> {
self.get_threads_usage(&[thread_id])
.await?
.into_iter()
.find(|usage| usage.thread_id == thread_id)
.ok_or_else(|| {
RequestError::from(anyhow!("thread usage response omitted requested thread"))
})
}
/// Reads at most 100 distinct threads; omitted results are unavailable, not zero.
pub async fn get_threads_usage(
&self,
thread_ids: &[&str],
) -> Result<Vec<ThreadUsage>, RequestError> {
let requested: std::collections::HashSet<_> = thread_ids.iter().copied().collect();
if thread_ids.is_empty() || thread_ids.len() > 100 || requested.len() != thread_ids.len() {
return Err(RequestError::from(anyhow!(
"expected 1–100 distinct thread IDs"
)));
}
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 });
let (body, _) = self.exec_request_detailed(request, "POST", &url).await?;
let response: ThreadUsageQueryResponse = serde_json::from_str(&body)
.map_err(|_| RequestError::Other(anyhow!("Invalid thread usage response.")))?;
let mut seen = std::collections::HashSet::new();
if response
.threads
.iter()
.any(|row| !requested.contains(row.thread_id.as_str()) || !seen.insert(&row.thread_id))
{
return Err(RequestError::from(anyhow!(
"thread usage returned unexpected threads"
)));
}
Ok(response
.threads
.into_iter()
.filter_map(|row| {
Some(ThreadUsage {
thread_id: row.thread_id,
estimated_usage_credits_micros: row.estimated_usage_credits_micros?,
estimated_usage_usd_micros: row.estimated_usage_usd_micros,
groups: row.groups.unwrap_or_default(),
})
})
.collect())
}
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;
|