fix: store OpenAI agent Generation span attributes under correct SpanAttributeKey constants by RudraDudhat2509 · Pull Request #24290 · mlflow/mlflow

RudraDudhat2509 · 2026-07-06T07:14:03Z

Related Issues/PRs

What changes are proposed in this pull request?

MlflowOpenAgentTracingProcessor._parse_span_data() wrote token usage for GENERATION spans under the plain key "usage". However, aggregate_usage_from_spans() and calculate_span_cost() both read SpanAttributeKey.CHAT_USAGE ("mlflow.chat.tokenUsage"). The key mismatch silently dropped all token counts from trace.info.token_usage. Streaming calls were hit hardest since the Agents SDK is the only source of usage data for those spans.

Two fixes in this PR:

Added _parse_generation_usage() helper that maps the Agents SDK usage dict to TokenUsageKey-keyed values and sets the result under SpanAttributeKey.CHAT_USAGE inside _parse_span_data(). Matches the existing pattern in mlflow/groq/_groq_autolog.py and mlflow/anthropic/autolog.py.
_parse_span_data() now sets SpanAttributeKey.MODEL (not plain "model") for GENERATION spans, so calculate_span_cost() can read the model name and compute cost correctly.

No API changes are needed.

How is this PR tested?

Existing unit/integration tests
New unit/integration tests
Manual tests

test_generation_span_attributes_stored_under_span_attribute_keys: mocks a GenerationSpanData, calls _parse_span_data(), asserts attributes contain both SpanAttributeKey.MODEL and SpanAttributeKey.CHAT_USAGE with correct values, and that neither plain "model" nor "usage" keys exist.

test_parse_generation_usage: parametrized over full dict, partial dict, nested input_tokens_details with both cache keys, partial cache details (only one key present), None, and empty dict.

test_calculate_span_cost_uses_generation_span_model_attribute: verifies cost calculation reads the model from SpanAttributeKey.MODEL and not the plain "model" key.

Does this PR require documentation update?

Does this PR require updating the MLflow Skills repository?

No.
Yes. Please link the corresponding PR or explain how you plan to update it.

Release Notes

Is this a user-facing change?

No.
Yes. Give a description of this change to be included in the release notes for MLflow users.

Token usage from OpenAI Agents SDK GENERATION spans (including streamed runs) now correctly rolls up into trace.info.token_usage and is used by calculate_span_cost(). Cached token counts (read and creation) are now captured from the nested input_tokens_details field. Model name is now stored under the correct attribute key so span cost is no longer silently skipped.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

How should the PR be classified in the release notes? Choose one:

rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
rn/feature - A new user-facing feature worth mentioning in the release notes
rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
rn/documentation - A user-facing documentation change worth mentioning in the release notes

Is this PR a critical bugfix or security fix that should go into the next patch release?

What is a minor/patch release?

Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
Minor releases are expected to contain larger changes, such as new features and improvements. Non-critical bug fixes and doc updates can be included as well. By default, your PR should target the next minor release.
Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
Patch releases are typically only performed when there has been a major regression or bug in the latest release. For the sake of stability, your PR should not be included in a patch release unless it is a critical fix, or if the risk level of your PR is exceedingly low.

This PR is critical and needs to be in the next patch release
This PR can wait for the next minor release

…teKey.CHAT_USAGE Signed-off-by: Rudra Dudhat <[email protected]>

github-actions · 2026-07-06T07:14:17Z

Install mlflow from this PR

# mlflow
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/24290/merge
# mlflow-skinny
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/24290/merge#subdirectory=libs/skinny

For Databricks, use the following command:

%sh curl -LsSf https://raw.githubusercontent.com/mlflow/mlflow/HEAD/dev/install-skinny.sh | sh -s pull/24290/merge

PR author's recent activity

In the last 14 days, @RudraDudhat2509 opened 15 PRs across 4 repos:

Repository	Open	Merged	Total
RudraDudhat2509/brok	0	7	7
RudraDudhat2509/mimic-eval	1	5	6
mlflow/mlflow	1	0	1
langfuse/langfuse-python	1	0	1

Copilot

Pull request overview

This PR fixes a data-correctness issue in MLflow’s OpenAI Agents SDK tracing integration where token usage recorded on GENERATION spans was stored under the wrong attribute key, preventing trace-level token aggregation (and related downstream consumers) from seeing those tokens.

Changes:

Update _parse_span_data() for GENERATION spans to store usage under SpanAttributeKey.CHAT_USAGE instead of a plain "usage" attribute.
Add _parse_generation_usage() helper to normalize usage dicts into MLflow’s TokenUsageKey schema.
Add unit tests covering both the span attribute behavior and the usage parsing helper.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
`mlflow/openai/_agent_tracer.py`	Writes Generation span usage under `SpanAttributeKey.CHAT_USAGE` via a new parsing helper.
`tests/openai/test_openai_agent_autolog.py`	Adds tests ensuring usage is stored under `CHAT_USAGE` and validating `_parse_generation_usage()`.

Co-authored-by: Copilot Autofix powered by AI <[email protected]> Signed-off-by: Rudra Dudhat <[email protected]>

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

RudraDudhat2509 · 2026-07-12T09:37:58Z

opened #24393 to address copilot comments

PattaraS · 2026-07-13T05:58:47Z

Following @joshuawong-db's note on #24393 (closed into this PR): could you also fold in the model-key fix here? The GENERATION span sets a plain "model" key, but calculate_span_cost() reads SpanAttributeKey.MODEL ("mlflow.llm.model"), so cost stays None for these traces even with the token-usage fix. Setting SpanAttributeKey.MODEL alongside the current change (plus a small test asserting cost/model resolves) would let this PR close both. The token-usage part already looks good.

…MODEL Using a plain string key "model" instead of SpanAttributeKey.MODEL ("mlflow.llm.model") caused calculate_span_cost() to silently skip cost calculation for Generation spans because the attribute was never found via get_attribute(). Signed-off-by: Rudra Dudhat <[email protected]>

RudraDudhat2509 · 2026-07-13T06:32:59Z

hey @PattaraS ! pushed the changes, also updated the test file, lmk if anything else is required from my end, thx

Merged model-key and token-usage assertions into one test (test_generation_span_attributes_stored_under_span_attribute_keys) to verify both fixes together; mocked cost calculation in test_calculate_span_cost_uses_generation_span_model_attribute to avoid coupling to LiteLLM pricing data. Signed-off-by: Rudra Dudhat <[email protected]>

Signed-off-by: Rudra Dudhat <[email protected]>

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

joshuawong-db · 2026-07-13T06:46:53Z

/review-v2

github-actions · 2026-07-13T07:00:19Z

Documentation preview for d06cc26 is available at:

https://pr-24290--mlflow-docs-preview.netlify.app/docs/latest/

More info

Ignore this comment if this PR does not change the documentation.
The preview is updated when a new commit is pushed to this PR.
This comment was created by this workflow run.
The documentation was built by this workflow run.

joshuawong-db

Thank you, in general this looks good, pending one comment on cached tokens nesting. Could you also update the PR description? It references tests that do not exist in the code.

…tion span usage Signed-off-by: Rudra Dudhat <[email protected]>

joshuawong-db

LGTM, thanks for the contribution @RudraDudhat2509

RudraDudhat2509 · 2026-07-14T07:43:50Z

Appreciate the quick review thx!

…AttributeKey constants (#24290) Signed-off-by: Rudra Dudhat <[email protected]> Co-authored-by: Copilot Autofix powered by AI <[email protected]> Co-authored-by: mprahl <[email protected]>

fix: store OpenAI agent Generation span token usage under SpanAttribu…

cc33814

…teKey.CHAT_USAGE Signed-off-by: Rudra Dudhat <[email protected]>

Copilot AI review requested due to automatic review settings July 6, 2026 07:14

github-actions Bot added the community Community/external contribution label Jul 6, 2026

github-actions Bot added size/M area/tracing MLflow Tracing and its integrations rn/bug-fix Mention under Bug Fixes in Changelogs. labels Jul 6, 2026

Copilot started reviewing on behalf of RudraDudhat2509 July 6, 2026 07:14 View session

Copilot AI reviewed Jul 6, 2026

View reviewed changes

Comment thread mlflow/openai/_agent_tracer.py

Comment thread mlflow/openai/_agent_tracer.py Outdated

fix: commit copilot suggestion

37dc072

Co-authored-by: Copilot Autofix powered by AI <[email protected]> Signed-off-by: Rudra Dudhat <[email protected]>

RudraDudhat2509 requested a review from Copilot July 11, 2026 10:55

Copilot started reviewing on behalf of RudraDudhat2509 July 11, 2026 10:55 View session

Copilot AI reviewed Jul 11, 2026

View reviewed changes

joshuawong-db mentioned this pull request Jul 13, 2026

[BUG] OpenAI Agents SDK GENERATION spans set plain "model" key, breaking calculate_span_cost() #24393

Closed

RudraDudhat2509 added 2 commits July 13, 2026 12:05

style: remove extra blank line between test functions

0fe7b5d

Signed-off-by: Rudra Dudhat <[email protected]>

RudraDudhat2509 requested a review from Copilot July 13, 2026 06:39

Copilot started reviewing on behalf of RudraDudhat2509 July 13, 2026 06:40 View session

Copilot AI reviewed Jul 13, 2026

View reviewed changes

github-actions Bot assigned joshuawong-db and PattaraS Jul 13, 2026

joshuawong-db reviewed Jul 13, 2026

View reviewed changes

Comment thread tests/openai/test_openai_agent_autolog.py

Comment thread mlflow/openai/_agent_tracer.py

fix: extract cached tokens from nested input_tokens_details in genera…

d06cc26

…tion span usage Signed-off-by: Rudra Dudhat <[email protected]>

RudraDudhat2509 changed the title ~~fix: store OpenAI agent Generation span token usage under SpanAttributeKey.CHAT_USAGE~~ fix: store OpenAI agent Generation span attributes under correct SpanAttributeKey constants Jul 13, 2026

RudraDudhat2509 requested a review from joshuawong-db July 13, 2026 10:58

joshuawong-db approved these changes Jul 14, 2026

View reviewed changes

joshuawong-db added this pull request to the merge queue Jul 14, 2026

Merged via the queue into mlflow:master with commit 43c1455 Jul 14, 2026
79 of 83 checks passed

RudraDudhat2509 deleted the fix/openai-agent-generation-span-token-usage branch July 14, 2026 11:47

Uh oh!

Conversation

RudraDudhat2509 commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Related Issues/PRs

What changes are proposed in this pull request?

How is this PR tested?

Does this PR require documentation update?

Does this PR require updating the MLflow Skills repository?

Release Notes

Is this a user-facing change?

What component(s), interfaces, languages, and integrations does this PR affect?

How should the PR be classified in the release notes? Choose one:

Is this PR a critical bugfix or security fix that should go into the next patch release?

Uh oh!

github-actions Bot commented Jul 6, 2026

Install mlflow from this PR

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

RudraDudhat2509 commented Jul 12, 2026

Uh oh!

PattaraS commented Jul 13, 2026

Uh oh!

RudraDudhat2509 commented Jul 13, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

joshuawong-db commented Jul 13, 2026

Uh oh!

github-actions Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

joshuawong-db left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

joshuawong-db left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

RudraDudhat2509 commented Jul 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

RudraDudhat2509 commented Jul 6, 2026 •

edited

Loading

github-actions Bot commented Jul 13, 2026 •

edited

Loading