feat(usage): attribute token usage [2 of 4]#24122
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b46cffa0fe
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| fn usage_contributors(&self) -> Vec<UsageContributor> { | ||
| Vec::new() | ||
| } |
There was a problem hiding this comment.
When a contributor-bearing handler is wrapped by override_tool_exposure (for example deferred_mcp_tools are registered this way), the wrapper falls back to this new default implementation because ExposureOverride does not implement usage_contributors(). In that context, deferred app/MCP/plugin tools return an empty contributor list, so their specs and subsequent tool outputs are never attributed even though they are the main source this feature is intended to report. Forward self.handler.usage_contributors() from the wrapper.
Useful? React with 👍 / 👎.
| 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, | ||
| } | ||
| })); |
There was a problem hiding this comment.
This records contributor token estimates before the later ToolSpec::Namespace filter is applied. When the provider lacks namespace_tools support (e.g. the Bedrock path covered in the spec-plan tests), direct MCP namespace specs are removed from model_visible_specs, but their contributors remain in router.usage_contributors(), so apps/MCP servers get charged for tool tokens that were never sent to the model and the denominator excludes those removed specs. Build attribution from the final visible specs, or drop contributors for filtered specs.
Useful? React with 👍 / 👎.
| @@ -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) | |||
There was a problem hiding this comment.
Because the contributor lookup table is populated only inside the exposure.is_direct() branch, tools exposed through the search/deferred path are never entered in usage_contributors_by_tool_name. When a deferred MCP/app/plugin tool is later loaded by tool_search and called, the following prompt contains its FunctionCall/FunctionCallOutput, but tool_result_contributors() cannot find provenance for that tool name, so the often-large returned content is not attributed to the app/MCP/plugin at all. Populate the call-result lookup for registered deferred tools even though their specs are not in the initial prompt.
Useful? React with 👍 / 👎.
| 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); |
There was a problem hiding this comment.
For multiple tools from the same MCP namespace/app, each runtime's spec() is a one-tool Namespace, so this estimates and later aggregates the namespace wrapper/description once per tool. The prompt denominator is computed from router.model_visible_specs() after merge_into_namespaces() has collapsed that wrapper to a single namespace, which means a connector with several tools can have source_estimated_tokens larger than the actual prompt tokens it contributed and receive inflated attribution. Derive contributor estimates from the final merged specs instead.
Useful? React with 👍 / 👎.
| usage_contributors.extend(runtime_usage_contributors.into_iter().map(|contributor| { | ||
| crate::usage::UsagePromptContributor { | ||
| contributor, | ||
| source_estimated_tokens: estimated_tokens, | ||
| } | ||
| })); |
There was a problem hiding this comment.
In code-mode-only sessions, app/MCP tools are hidden as direct tools and represented inside the generated exec tool description, but this loop only asks the top-level runtime for contributors. The CodeModeExecuteHandler uses the default empty contributor list, so the app/MCP/plugin tool definitions embedded in the model-visible exec spec are recorded as unattributed even though those tokens were sent to the model. Carry the nested runtimes' contributors into the exec handler or attribute the generated exec spec back to those nested tools.
Useful? React with 👍 / 👎.
fcoury-oai
changed the title
fcoury-oai
force-pushed
the
fcoury/usage-storage
branch
from
c3ef233 to
f5afec4
Compare
fcoury-oai
force-pushed
the
fcoury/usage-attribution
branch
from
b46cffa to
0d9682b
Compare
fcoury-oai
force-pushed
the
fcoury/usage-storage
branch
from
f5afec4 to
357d711
Compare
fcoury-oai
force-pushed
the
fcoury/usage-attribution
branch
from
0d9682b to
6f093f8
Compare
fcoury-oai
force-pushed
the
fcoury/usage-attribution
branch
from
3ff5b6b to
219664b
Compare
1 participant
Why
The storage layer needs real contributor samples before any report can be useful. This PR records the prompt/tool provenance needed to explain which local features consumed tokens.
What Changed
How to Test
Targeted tests:
cargo test -p codex-core usageStack
/usagecommand