From 83505d85cfc184c6d2f407c9e0861ebf232b1cf8 Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Fri, 29 May 2026 12:38:19 -0300 Subject: [PATCH] refactor(usage): borrow prompt attribution inputs --- codex-rs/core/src/session/turn.rs | 4 +- codex-rs/core/src/tools/router.rs | 9 ++- codex-rs/core/src/tools/spec_plan_tests.rs | 8 +-- codex-rs/core/src/usage.rs | 84 ++++++++++------------ 4 files changed, 49 insertions(+), 56 deletions(-) diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 14b5769dbf..cc0b4c2414 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -900,14 +900,16 @@ pub(crate) fn build_prompt( turn_context: &TurnContext, base_instructions: BaseInstructions, ) -> Prompt { + let tools = router.model_visible_specs(); let usage_attribution = crate::usage::UsagePromptAttribution::from_prompt( &input, + &tools, router, base_instructions.text.as_str(), ); Prompt { input, - tools: router.model_visible_specs(), + tools, parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls, usage_attribution, base_instructions, diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 1a93fcaa89..ffc124ade0 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -74,18 +74,17 @@ impl ToolRouter { self.model_visible_specs.clone() } - pub(crate) fn usage_contributors(&self) -> Vec { - self.usage_contributors.clone() + pub(crate) fn usage_contributors(&self) -> &[crate::usage::UsagePromptContributor] { + &self.usage_contributors } 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() + .map_or(&[], Vec::as_slice) } #[cfg(test)] diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index dd755f74d2..0c70f086c3 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -88,8 +88,8 @@ impl ToolPlanProbe { .collect::>(); let prompt_usage_contributor_labels = router .usage_contributors() - .into_iter() - .map(|contributor| contributor.contributor.label) + .iter() + .map(|contributor| contributor.contributor.label.clone()) .collect::>(); let tool_usage_contributor_labels = registered_tool_names .iter() @@ -99,8 +99,8 @@ impl ToolPlanProbe { ( name.to_string(), contributors - .into_iter() - .map(|contributor| contributor.label) + .iter() + .map(|contributor| contributor.label.clone()) .collect::>(), ) }) diff --git a/codex-rs/core/src/usage.rs b/codex-rs/core/src/usage.rs index df6ec1d87a..26fb316ce6 100644 --- a/codex-rs/core/src/usage.rs +++ b/codex-rs/core/src/usage.rs @@ -7,6 +7,7 @@ use codex_protocol::protocol::UsageAttributionItem; use codex_protocol::protocol::UsageContributor; use codex_protocol::protocol::UsageContributorKind; use codex_tools::ToolName; +use codex_tools::ToolSpec; use codex_utils_output_truncation::approx_token_count; use std::collections::BTreeMap; use std::collections::HashMap; @@ -26,18 +27,18 @@ pub(crate) struct UsagePromptContributor { impl UsagePromptAttribution { pub(crate) fn from_prompt( input: &[ResponseItem], + tools: &[ToolSpec], router: &ToolRouter, base_instructions: &str, ) -> Self { let mut contributors = skill_contributors(input); - contributors.extend(router.usage_contributors()); + contributors.extend_from_slice(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() + let tool_tokens = tools .iter() .map(estimate_serialized_tokens) .fold(0i64, i64::saturating_add); @@ -105,56 +106,47 @@ fn tool_result_contributors( ) -> Vec { let contributors_by_call_id = input .iter() - .filter_map(|item| tool_call_contributors(item, router)) + .filter_map(|item| { + 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)) + }) .collect::>(); input .iter() - .filter_map(|item| tool_result_contributor(item, &contributors_by_call_id)) + .filter_map(|item| { + 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(move |contributor| UsagePromptContributor { + contributor, + source_estimated_tokens, + }), + ) + }) .flatten() .collect() } -fn tool_call_contributors( - item: &ResponseItem, - router: &ToolRouter, -) -> Option<(String, Vec)> { - 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>, -) -> Option> { - 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 { let ResponseItem::Message { content, .. } = item else { return None;