openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 17:38:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(usage): attribute token usage

Browse Source
This commit is contained in:
Felipe Coury
2026-05-28 21:05:18 -03:00
parent da401bd943
commit 43d4c2ee73
12 changed files with 522 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/client_common.rs
+4
View File
@@ -1,3 +1,4 @@
use crate::usage::UsagePromptAttribution;
pub use codex_api::ResponseEvent;
use codex_config::types::Personality;
use codex_protocol::error::Result;
@@ -33,6 +34,8 @@ pub struct Prompt {
/// Whether parallel tool calls are permitted for this prompt.
pub(crate) parallel_tool_calls: bool,
pub(crate) usage_attribution: UsagePromptAttribution,
pub base_instructions: BaseInstructions,
/// Optionally specify the personality of the model.
@@ -51,6 +54,7 @@ impl Default for Prompt {
input: Vec::new(),
tools: Vec::new(),
parallel_tool_calls: false,
usage_attribution: UsagePromptAttribution::default(),
base_instructions: BaseInstructions::default(),
personality: None,
output_schema: None,
codex-rs/core/src/compact_remote.rs
+1
View File
@@ -194,6 +194,7 @@ async fn run_remote_compact_task_inner_impl(
input: prompt_input,
tools: tool_router.model_visible_specs(),
parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls,
usage_attribution: Default::default(),
base_instructions,
personality: turn_context.personality,
output_schema: None,
codex-rs/core/src/compact_remote_v2.rs
+1
View File
@@ -211,6 +211,7 @@ async fn run_remote_compact_task_inner_impl(
input,
tools: tool_router.model_visible_specs(),
parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls,
usage_attribution: Default::default(),
base_instructions,
personality: turn_context.personality,
output_schema: None,
codex-rs/core/src/lib.rs
+1
View File
@@ -99,6 +99,7 @@ pub(crate) use skills::skills_load_input_from_config;
mod stream_events_utils;
pub mod test_support;
mod unified_exec;
mod usage;
pub mod windows_sandbox;
pub use client::X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER;
pub use codex_protocol::config_types::ModelProviderAuthInfo;
codex-rs/core/src/session/mod.rs
+30
View File
@@ -2947,6 +2947,36 @@ impl Session {
}
}
pub(crate) async fn record_usage_attribution(
&self,
turn_context: &TurnContext,
prompt: &crate::client_common::Prompt,
response_id: &str,
token_usage: Option<&TokenUsage>,
) {
let Some(token_usage) = token_usage else {
return;
};
let occurred_at = chrono::Utc::now().timestamp();
let attribution = prompt.usage_attribution.complete(
format!("{}:{response_id}", self.conversation_id),
turn_context.sub_id.clone(),
response_id.to_string(),
occurred_at,
token_usage.clone(),
);
if let Some(state_db) = self.state_db()
&& let Err(err) = state_db
.record_usage_sample(&codex_state::UsageSample {
thread_id: self.conversation_id,
attribution,
})
.await
{
warn!("failed to persist usage sample: {err}");
}
}
pub(crate) async fn recompute_token_usage(&self, turn_context: &TurnContext) {
let history = self.clone_history().await;
let base_instructions = self.get_base_instructions().await;
codex-rs/core/src/session/turn.rs
+13
View File
@@ -900,10 +900,16 @@ pub(crate) fn build_prompt(
turn_context: &TurnContext,
base_instructions: BaseInstructions,
) -> Prompt {
let usage_attribution = crate::usage::UsagePromptAttribution::from_prompt(
&input,
router,
base_instructions.text.as_str(),
);
Prompt {
input,
tools: router.model_visible_specs(),
parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls,
usage_attribution,
base_instructions,
personality: turn_context.personality,
output_schema: turn_context.final_output_json_schema.clone(),
@@ -2006,6 +2012,13 @@ async fn try_run_sampling_request(
&mut assistant_message_stream_parsers,
)
.await;
sess.record_usage_attribution(
&turn_context,
prompt,
response_id.as_str(),
token_usage.as_ref(),
)
.await;
sess.record_token_usage_info(&turn_context, token_usage.as_ref())
.await;
should_emit_token_count = true;
codex-rs/core/src/tools/handlers/mcp.rs
+34
View File
@@ -17,6 +17,8 @@ use crate::tools::registry::ToolExecutor;
use crate::tools::registry::ToolTelemetryTags;
use crate::tools::tool_search_entry::ToolSearchInfo;
use codex_mcp::ToolInfo;
use codex_protocol::protocol::UsageContributor;
use codex_protocol::protocol::UsageContributorKind;
use codex_tools::ResponsesApiNamespace;
use codex_tools::ResponsesApiNamespaceTool;
use codex_tools::ToolName;
@@ -131,6 +133,38 @@ impl ToolExecutor<ToolInvocation> for McpHandler {
}
impl CoreToolRuntime for McpHandler {
fn usage_contributors(&self) -> Vec<UsageContributor> {
let mut contributors = Vec::new();
if let Some(connector_id) = self.tool_info.connector_id.as_ref() {
contributors.push(UsageContributor {
kind: UsageContributorKind::App,
id: connector_id.clone(),
label: self
.tool_info
.connector_name
.clone()
.unwrap_or_else(|| connector_id.clone()),
});
} else {
contributors.push(UsageContributor {
kind: UsageContributorKind::McpServer,
id: self.tool_info.server_name.clone(),
label: self.tool_info.server_name.clone(),
});
}
contributors.extend(
self.tool_info
.plugin_display_names
.iter()
.map(|plugin_name| UsageContributor {
kind: UsageContributorKind::Plugin,
id: plugin_name.clone(),
label: plugin_name.clone(),
}),
);
contributors
}
fn search_info(&self) -> Option<ToolSearchInfo> {
let source_name = self
.tool_info
codex-rs/core/src/tools/parallel.rs
+2
View File
@@ -423,6 +423,8 @@ mod tests {
let router = Arc::new(ToolRouter::from_parts(
ToolRegistry::from_tools([handler]),
Vec::new(),
Vec::new(),
std::collections::HashMap::new(),
));
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let runtime = ToolCallRuntime::new(router, session, turn_context, tracker);
codex-rs/core/src/tools/registry.rs
+5
View File
@@ -30,6 +30,7 @@ use codex_extension_api::ToolCallOutcome;
use codex_protocol::models::FunctionCallOutputPayload;
use codex_protocol::models::ResponseInputItem;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::UsageContributor;
use codex_tools::ToolName;
use codex_tools::ToolSpec;
use futures::future::BoxFuture;
@@ -104,6 +105,10 @@ pub(crate) trait CoreToolRuntime: ToolExecutor<ToolInvocation> {
})
}
fn usage_contributors(&self) -> Vec<UsageContributor> {
Vec::new()
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
let ToolPayload::Function { arguments } = &invocation.payload else {
return None;
codex-rs/core/src/tools/router.rs
+29 -1
View File
@@ -17,6 +17,7 @@ use codex_tools::ToolCall as ExtensionToolCall;
use codex_tools::ToolExecutor;
use codex_tools::ToolName;
use codex_tools::ToolSpec;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use tokio_util::sync::CancellationToken;
@@ -34,6 +35,9 @@ pub struct ToolCall {
pub struct ToolRouter {
registry: ToolRegistry,
model_visible_specs: Vec<ToolSpec>,
usage_contributors: Vec<crate::usage::UsagePromptContributor>,
usage_contributors_by_tool_name:
HashMap<ToolName, Vec<codex_protocol::protocol::UsageContributor>>,
}
pub(crate) struct ToolRouterParams<'a> {
@@ -49,10 +53,20 @@ impl ToolRouter {
build_tool_router(turn_context, params)
}
pub(crate) fn from_parts(registry: ToolRegistry, model_visible_specs: Vec<ToolSpec>) -> Self {
pub(crate) fn from_parts(
registry: ToolRegistry,
model_visible_specs: Vec<ToolSpec>,
usage_contributors: Vec<crate::usage::UsagePromptContributor>,
usage_contributors_by_tool_name: HashMap<
ToolName,
Vec<codex_protocol::protocol::UsageContributor>,
>,
) -> Self {
Self {
registry,
model_visible_specs,
usage_contributors,
usage_contributors_by_tool_name,
}
}
@@ -60,6 +74,20 @@ impl ToolRouter {
self.model_visible_specs.clone()
}
pub(crate) fn usage_contributors(&self) -> Vec<crate::usage::UsagePromptContributor> {
self.usage_contributors.clone()
}
pub(crate) fn usage_contributors_for_tool_name(
&self,
tool_name: &ToolName,
) -> Vec<codex_protocol::protocol::UsageContributor> {
self.usage_contributors_by_tool_name
.get(tool_name)
.cloned()
.unwrap_or_default()
}
#[cfg(test)]
pub(crate) fn registered_tool_names_for_test(&self) -> Vec<ToolName> {
self.registry.tool_names_for_test()
codex-rs/core/src/tools/spec_plan.rs
+39 -5
View File
@@ -80,6 +80,7 @@ use codex_tools::request_user_input_available_modes;
use codex_tools::shell_command_backend_for_features;
use codex_tools::shell_type_for_model_and_features;
use std::collections::BTreeMap;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use tracing::warn;
@@ -146,14 +147,25 @@ pub(crate) fn build_tool_router(
turn_context: &TurnContext,
params: ToolRouterParams<'_>,
) -> ToolRouter {
let (model_visible_specs, registry) = build_tool_specs_and_registry(turn_context, params);
ToolRouter::from_parts(registry, model_visible_specs)
let (model_visible_specs, registry, usage_contributors, usage_contributors_by_tool_name) =
build_tool_specs_and_registry(turn_context, params);
ToolRouter::from_parts(
registry,
model_visible_specs,
usage_contributors,
usage_contributors_by_tool_name,
)
}
fn build_tool_specs_and_registry(
turn_context: &TurnContext,
params: ToolRouterParams<'_>,
) -> (Vec<ToolSpec>, ToolRegistry) {
) -> (
Vec<ToolSpec>,
ToolRegistry,
Vec<crate::usage::UsagePromptContributor>,
HashMap<ToolName, Vec<codex_protocol::protocol::UsageContributor>>,
) {
let ToolRouterParams {
mcp_tools,
deferred_mcp_tools,
@@ -183,12 +195,19 @@ fn build_tool_specs_and_registry(
fn build_model_visible_specs_and_registry(
turn_context: &TurnContext,
planned_tools: PlannedTools,
) -> (Vec<ToolSpec>, ToolRegistry) {
) -> (
Vec<ToolSpec>,
ToolRegistry,
Vec<crate::usage::UsagePromptContributor>,
HashMap<ToolName, Vec<codex_protocol::protocol::UsageContributor>>,
) {
let PlannedTools {
runtimes,
hosted_specs,
} = planned_tools;
let mut specs = Vec::new();
let mut usage_contributors = Vec::new();
let mut usage_contributors_by_tool_name = HashMap::new();
let mut seen_tool_names = HashSet::new();
for runtime in &runtimes {
let tool_name = runtime.tool_name();
@@ -199,6 +218,16 @@ fn build_model_visible_specs_and_registry(
if exposure.is_direct() && !is_hidden_by_code_mode_only(turn_context, &tool_name, exposure)
{
let spec = runtime.spec();
let estimated_tokens = crate::usage::estimate_serialized_tokens(&spec);
let runtime_usage_contributors = runtime.usage_contributors();
usage_contributors_by_tool_name
.insert(tool_name.clone(), runtime_usage_contributors.clone());
usage_contributors.extend(runtime_usage_contributors.into_iter().map(|contributor| {
crate::usage::UsagePromptContributor {
contributor,
source_estimated_tokens: estimated_tokens,
}
}));
specs.push(spec_for_model_request(turn_context, exposure, spec));
}
}
@@ -220,7 +249,12 @@ fn build_model_visible_specs_and_registry(
})
.collect();
(model_visible_specs, registry)
(
model_visible_specs,
registry,
usage_contributors,
usage_contributors_by_tool_name,
)
}
fn spec_for_model_request(
codex-rs/core/src/usage.rs
+363
View File
@@ -0,0 +1,363 @@
use crate::tools::router::ToolRouter;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::TokenUsage;
use codex_protocol::protocol::UsageAttributionContributor;
use codex_protocol::protocol::UsageAttributionItem;
use codex_protocol::protocol::UsageContributor;
use codex_protocol::protocol::UsageContributorKind;
use codex_tools::ToolName;
use codex_utils_output_truncation::approx_token_count;
use std::collections::BTreeMap;
use std::collections::HashMap;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct UsagePromptAttribution {
pub(crate) prompt_estimated_tokens: i64,
pub(crate) contributors: Vec<UsagePromptContributor>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct UsagePromptContributor {
pub(crate) contributor: UsageContributor,
pub(crate) source_estimated_tokens: i64,
}
impl UsagePromptAttribution {
pub(crate) fn from_prompt(
input: &[ResponseItem],
router: &ToolRouter,
base_instructions: &str,
) -> Self {
let mut contributors = skill_contributors(input);
contributors.extend(router.usage_contributors());
contributors.extend(tool_result_contributors(input, router));
let input_tokens = input
.iter()
.map(estimate_response_item_tokens)
.fold(0i64, i64::saturating_add);
let tool_tokens = router
.model_visible_specs()
.iter()
.map(estimate_serialized_tokens)
.fold(0i64, i64::saturating_add);
let base_tokens = i64::try_from(approx_token_count(base_instructions)).unwrap_or(i64::MAX);
Self {
prompt_estimated_tokens: base_tokens
.saturating_add(input_tokens)
.saturating_add(tool_tokens),
contributors: aggregate_contributors(contributors),
}
}
pub(crate) fn complete(
&self,
sample_id: String,
turn_id: String,
response_id: String,
occurred_at: i64,
token_usage: TokenUsage,
) -> UsageAttributionItem {
let non_cached_input = token_usage.non_cached_input();
let contributors = self
.contributors
.iter()
.map(|contributor| UsageAttributionContributor {
contributor: contributor.contributor.clone(),
source_estimated_tokens: contributor.source_estimated_tokens,
attributed_tokens: attributable_tokens(
non_cached_input,
contributor.source_estimated_tokens,
self.prompt_estimated_tokens,
),
})
.filter(|contributor| contributor.attributed_tokens > 0)
.collect();
UsageAttributionItem {
sample_id,
turn_id,
response_id,
occurred_at,
token_usage,
prompt_estimated_tokens: self.prompt_estimated_tokens,
contributors,
}
}
}
pub(crate) fn estimate_serialized_tokens<T: serde::Serialize>(value: &T) -> i64 {
serde_json::to_string(value)
.map(|serialized| i64::try_from(approx_token_count(&serialized)).unwrap_or(i64::MAX))
.unwrap_or(/*default*/ 0)
}
fn estimate_response_item_tokens(item: &ResponseItem) -> i64 {
estimate_serialized_tokens(item)
}
fn skill_contributors(input: &[ResponseItem]) -> Vec<UsagePromptContributor> {
input.iter().filter_map(skill_contributor).collect()
}
fn tool_result_contributors(
input: &[ResponseItem],
router: &ToolRouter,
) -> Vec<UsagePromptContributor> {
let contributors_by_call_id = input
.iter()
.filter_map(|item| tool_call_contributors(item, router))
.collect::<HashMap<_, _>>();
input
.iter()
.filter_map(|item| tool_result_contributor(item, &contributors_by_call_id))
.flatten()
.collect()
}
fn tool_call_contributors(
item: &ResponseItem,
router: &ToolRouter,
) -> Option<(String, Vec<UsageContributor>)> {
let (call_id, tool_name) = match item {
ResponseItem::FunctionCall {
call_id,
name,
namespace,
..
} => (call_id, ToolName::new(namespace.clone(), name)),
ResponseItem::CustomToolCall { call_id, name, .. } => (call_id, ToolName::plain(name)),
_ => return None,
};
let contributors = router.usage_contributors_for_tool_name(&tool_name);
(!contributors.is_empty()).then(|| (call_id.clone(), contributors))
}
fn tool_result_contributor(
item: &ResponseItem,
contributors_by_call_id: &HashMap<String, Vec<UsageContributor>>,
) -> Option<Vec<UsagePromptContributor>> {
let call_id = match item {
ResponseItem::FunctionCallOutput { call_id, .. }
| ResponseItem::CustomToolCallOutput { call_id, .. } => call_id,
_ => return None,
};
let source_estimated_tokens = estimate_response_item_tokens(item);
Some(
contributors_by_call_id
.get(call_id)?
.iter()
.cloned()
.map(|contributor| UsagePromptContributor {
contributor,
source_estimated_tokens,
})
.collect(),
)
}
fn skill_contributor(item: &ResponseItem) -> Option<UsagePromptContributor> {
let ResponseItem::Message { content, .. } = item else {
return None;
};
let text = content.iter().find_map(|content| match content {
ContentItem::InputText { text } if text.contains("<skill>") => Some(text.as_str()),
_ => None,
})?;
let name = tag_contents(text, "name")?;
let path = tag_contents(text, "path")?;
Some(UsagePromptContributor {
contributor: UsageContributor {
kind: UsageContributorKind::Skill,
id: path.to_string(),
label: name.to_string(),
},
source_estimated_tokens: i64::try_from(approx_token_count(text)).unwrap_or(i64::MAX),
})
}
fn tag_contents<'a>(text: &'a str, tag: &str) -> Option<&'a str> {
let open = format!("<{tag}>");
let close = format!("</{tag}>");
let start = text.find(open.as_str())? + open.len();
let end = text[start..].find(close.as_str())? + start;
Some(text[start..end].trim())
}
fn aggregate_contributors(
contributors: Vec<UsagePromptContributor>,
) -> Vec<UsagePromptContributor> {
let mut aggregated = BTreeMap::new();
for contributor in contributors {
let key = (
contributor.contributor.kind as u8,
contributor.contributor.id.clone(),
contributor.contributor.label.clone(),
);
aggregated
.entry(key)
.and_modify(|existing: &mut UsagePromptContributor| {
existing.source_estimated_tokens = existing
.source_estimated_tokens
.saturating_add(contributor.source_estimated_tokens);
})
.or_insert(contributor);
}
aggregated.into_values().collect()
}
fn attributable_tokens(non_cached_input: i64, source_tokens: i64, prompt_tokens: i64) -> i64 {
if non_cached_input <= 0 || source_tokens <= 0 || prompt_tokens <= 0 {
return 0;
}
non_cached_input
.saturating_mul(source_tokens)
.saturating_add(prompt_tokens / 2)
/ prompt_tokens
}
#[cfg(test)]
mod tests {
use super::*;
use crate::tools::registry::ToolRegistry;
use codex_protocol::models::FunctionCallOutputPayload;
use pretty_assertions::assert_eq;
#[test]
fn complete_attributes_only_non_cached_input_tokens() {
let attribution = UsagePromptAttribution {
prompt_estimated_tokens: 100,
contributors: vec![
usage_prompt_contributor(
UsageContributorKind::Skill,
"/skills/tmux",
"tmux",
/*source_estimated_tokens*/ 25,
),
usage_prompt_contributor(
UsageContributorKind::App,
"slack",
"Slack",
/*source_estimated_tokens*/ 10,
),
],
};
let usage = attribution.complete(
"sample".to_string(),
"turn".to_string(),
"response".to_string(),
/*occurred_at*/ 1_700_000_000,
TokenUsage {
input_tokens: 100,
cached_input_tokens: 40,
output_tokens: 20,
reasoning_output_tokens: 0,
total_tokens: 120,
},
);
assert_eq!(
usage.contributors,
vec![
UsageAttributionContributor {
contributor: usage_contributor(
UsageContributorKind::Skill,
"/skills/tmux",
"tmux",
),
source_estimated_tokens: 25,
attributed_tokens: 15,
},
UsageAttributionContributor {
contributor: usage_contributor(UsageContributorKind::App, "slack", "Slack"),
source_estimated_tokens: 10,
attributed_tokens: 6,
},
]
);
}
#[test]
fn skill_contributors_use_skill_path_as_stable_id() {
let item = ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "<skill><name>tmux</name><path>/skills/tmux/SKILL.md</path></skill>"
.to_string(),
}],
phase: None,
};
assert_eq!(
skill_contributors(&[item]),
vec![UsagePromptContributor {
contributor: usage_contributor(
UsageContributorKind::Skill,
"/skills/tmux/SKILL.md",
"tmux",
),
source_estimated_tokens: i64::try_from(approx_token_count(
"<skill><name>tmux</name><path>/skills/tmux/SKILL.md</path></skill>",
))
.expect("skill prompt token estimate should fit in i64"),
}]
);
}
#[test]
fn tool_results_reuse_tool_usage_provenance() {
let contributor = usage_contributor(UsageContributorKind::App, "slack", "Slack");
let tool_name = ToolName::plain("mcp__slack__search");
let router = ToolRouter::from_parts(
ToolRegistry::from_tools(Vec::<
std::sync::Arc<dyn crate::tools::registry::CoreToolRuntime>,
>::new()),
Vec::new(),
Vec::new(),
HashMap::from([(tool_name.clone(), vec![contributor.clone()])]),
);
let tool_result = ResponseItem::FunctionCallOutput {
call_id: "call-1".to_string(),
output: FunctionCallOutputPayload::from_text("result".to_string()),
};
let input = vec![
ResponseItem::FunctionCall {
id: None,
name: tool_name.name,
namespace: tool_name.namespace,
arguments: "{}".to_string(),
call_id: "call-1".to_string(),
},
tool_result.clone(),
];
assert_eq!(
tool_result_contributors(&input, &router),
vec![UsagePromptContributor {
contributor,
source_estimated_tokens: estimate_response_item_tokens(&tool_result),
}]
);
}
fn usage_prompt_contributor(
kind: UsageContributorKind,
id: &str,
label: &str,
source_estimated_tokens: i64,
) -> UsagePromptContributor {
UsagePromptContributor {
contributor: usage_contributor(kind, id, label),
source_estimated_tokens,
}
}
fn usage_contributor(kind: UsageContributorKind, id: &str, label: &str) -> UsageContributor {
UsageContributor {
kind,
id: id.to_string(),
label: label.to_string(),
}
}
}