refactor(usage): borrow prompt attribution inputs

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,

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<crate::usage::UsagePromptContributor> {
-        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> {
+    ) -> &[codex_protocol::protocol::UsageContributor] {
         self.usage_contributors_by_tool_name
             .get(tool_name)
-            .cloned()
-            .unwrap_or_default()
+            .map_or(&[], Vec::as_slice)
     }
 
     #[cfg(test)]

codex-rs/core/src/tools/spec_plan_tests.rs

@@ -88,8 +88,8 @@ impl ToolPlanProbe {
             .collect::<Vec<_>>();
         let prompt_usage_contributor_labels = router
             .usage_contributors()
-            .into_iter()
-            .map(|contributor| contributor.contributor.label)
+            .iter()
+            .map(|contributor| contributor.contributor.label.clone())
             .collect::<Vec<_>>();
         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::<Vec<_>>(),
                     )
                 })

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<UsagePromptContributor> {
     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::<HashMap<_, _>>();
     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<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;