Lifei/fixed accumulated token count #6587

lifeizhou-ap · 2026-01-20T09:46:47Z

Summary

Fix wrong accumulated token count for databricks claude models.

Why
For Databricks Claude models, the usage is in all the chunks with same prompt_token. But only the last one with finish_reason has completion_token.

With the change in this PR #6493, for each api,
Actual: we adds prompt_token in all the chunks
Expected: we should only add once, which can use the chunk with finish_reason.

What

Only use the usage with completion_token. If completion_token is null, this means we don't have the usage yet.
Added usage in the tool_call_chunk too (I think this was an existing issue)
Added tests with api response test data

Type of Change

AI Assistance

This PR was created or reviewed with AI assistance

Testing

Unit tests and manual tests

Copilot

Pull request overview

This PR fixes an issue with accumulated token count in streaming responses from OpenAI-compatible providers. The changes extract and update usage information from streaming chunks, ensuring that usage data with output tokens is correctly captured and yielded only once.

Changes:

Added extract_usage_with_output_tokens helper function to filter usage data that includes output tokens
Modified streaming response handler to accumulate usage from chunks during tool call processing
Refactored tests with helper functions and added comprehensive test coverage for multiple providers (OpenRouter, OpenAI GPT-5, Tetrate Claude)

Copilot · 2026-01-20T09:48:43Z

crates/goose/src/providers/formats/openai.rs

                                let tool_chunk: StreamingChunk = serde_json::from_str(line)
                                    .map_err(|e| anyhow!("Failed to parse streaming chunk: {}: {:?}", e, &line))?;

+                                if let Some(chunk_usage) = extract_usage_witqh_output_tokens(&tool_chunk) {


Typo in function name: "extract_usage_witqh_output_tokens" should be "extract_usage_with_output_tokens".

Suggested change

if let Some(chunk_usage) = extract_usage_witqh_output_tokens(&tool_chunk) {

if let Some(chunk_usage) = extract_usage_with_output_tokens(&tool_chunk) {

Copilot · 2026-01-20T09:48:43Z

crates/goose/src/providers/formats/openai.rs

    #[tokio::test]
    async fn test_streamed_multi_tool_response_to_messages() -> anyhow::Result<()> {
-        let response_lines = r#"
+        let responqse_lines = r#"


Typo in variable name: "responqse_lines" should be "response_lines".

Suggested change

let responqse_lines = r#"

let response_lines = r#"

DOsinga · 2026-01-20T14:01:17Z

crates/goose/src/providers/formats/openai.rs

+        assert_eq!(usage.usage.total_tokens, Some(expected_total));
+    }
+
    #[tokio::test]


I haven't looked at the details since this just feels like a large dump, but not sure we need three tests here to verify that we are now doing the right thing?

The bug is caused by different structures of the streaming chunks in api response. So I added the tests with the test data with variations for different providers. This will provide harness when we changing the logic in the future.

jamadeo

Is it the same across all providers @lifeizhou-ap ? You mention databricks/claude but since this code is in formats/openai.rs I'd be worried if we're missing some specific logic per-provider for token tracking

lifeizhou-ap · 2026-01-20T23:25:48Z

Is it the same across all providers @lifeizhou-ap ? You mention databricks/claude but since this code is in formats/openai.rs I'd be worried if we're missing some specific logic per-provider for token tracking

Yes, I was worrying about this too including the last PR. I manually tested with different providers again. Also I added the chunk payload in new tests to make sure future change won't break this.

I am not sure whether I understand your question correctly, happy to have a chat!

…ased * 'main' of github.com:block/goose: Lifei/fixed accumulated token count (#6587) Dont show MCP UI/Apps until tool is approved (#6492)

…ovider * 'main' of github.com:block/goose: increase worker threads for ci (#6614) docs: todo tutorial update (#6613) Added goose doc map md file for goose agent to find relevant doc easily. (#6598) add back goose branding to home (#6617) fix: actually set the working dir for extensions from session (#6612) Multi chat (#6428) Lifei/fixed accumulated token count (#6587) Dont show MCP UI/Apps until tool is approved (#6492) docs: max tokens config (#6596) User configurable templates (#6420) docs: http proxy environment variables (#6594) feat: exclude subagent tool from code_execution filtering (#6531)

* main: increase worker threads for ci (#6614) docs: todo tutorial update (#6613) Added goose doc map md file for goose agent to find relevant doc easily. (#6598) add back goose branding to home (#6617) fix: actually set the working dir for extensions from session (#6612) Multi chat (#6428) Lifei/fixed accumulated token count (#6587) Dont show MCP UI/Apps until tool is approved (#6492) docs: max tokens config (#6596) User configurable templates (#6420) docs: http proxy environment variables (#6594) feat: exclude subagent tool from code_execution filtering (#6531) Fix path for global agent skills (#6591) recipes: add mcp server (#6552) feat(gcp-vertex): add model list with org policy filtering (#6393) chore: encourage extension searching (#6582) blog: mobile apps consolidation and roadmap (#6580) chore: remove unused dependencies in cargo.toml (#6561) resolved all the extensions to load in cli (#6464)

* main: increase worker threads for ci (#6614) docs: todo tutorial update (#6613) Added goose doc map md file for goose agent to find relevant doc easily. (#6598) add back goose branding to home (#6617) fix: actually set the working dir for extensions from session (#6612) Multi chat (#6428) Lifei/fixed accumulated token count (#6587) Dont show MCP UI/Apps until tool is approved (#6492)

* main: (41 commits) chore: tweak release docs (#6571) fix(goose): propagate session_id across providers and MCP (#6584) increase worker threads for ci (#6614) docs: todo tutorial update (#6613) Added goose doc map md file for goose agent to find relevant doc easily. (#6598) add back goose branding to home (#6617) fix: actually set the working dir for extensions from session (#6612) Multi chat (#6428) Lifei/fixed accumulated token count (#6587) Dont show MCP UI/Apps until tool is approved (#6492) docs: max tokens config (#6596) User configurable templates (#6420) docs: http proxy environment variables (#6594) feat: exclude subagent tool from code_execution filtering (#6531) Fix path for global agent skills (#6591) recipes: add mcp server (#6552) feat(gcp-vertex): add model list with org policy filtering (#6393) chore: encourage extension searching (#6582) blog: mobile apps consolidation and roadmap (#6580) chore: remove unused dependencies in cargo.toml (#6561) ...

* main: (68 commits) fix(docs): use dynamic import for globby ESM module (#6636) chore: trigger CI Document tab completion (#6635) Install goose-mcp crate dependencies (#6632) feat(goose): standardize agent-session-id for session correlation (#6626) chore: tweak release docs (#6571) fix(goose): propagate session_id across providers and MCP (#6584) increase worker threads for ci (#6614) docs: todo tutorial update (#6613) Added goose doc map md file for goose agent to find relevant doc easily. (#6598) add back goose branding to home (#6617) fix: actually set the working dir for extensions from session (#6612) Multi chat (#6428) Lifei/fixed accumulated token count (#6587) Dont show MCP UI/Apps until tool is approved (#6492) docs: max tokens config (#6596) User configurable templates (#6420) docs: http proxy environment variables (#6594) feat: exclude subagent tool from code_execution filtering (#6531) Fix path for global agent skills (#6591) ...

Signed-off-by: fbalicchia <fbalicchia@cuebiq.com>

lifeizhou-ap added 2 commits January 20, 2026 18:52

fixed the usage for databricks sonnet models

4751ab8

added test for streaming

de13baf

Copilot AI review requested due to automatic review settings January 20, 2026 09:46

Copilot started reviewing on behalf of lifeizhou-ap January 20, 2026 09:47 View session

Copilot AI reviewed Jan 20, 2026

View reviewed changes

fixed typo

506568b

lifeizhou-ap mentioned this pull request Jan 20, 2026

fixed 0 token in openrouter steaming #6493

Merged

10 tasks

lifeizhou-ap requested a review from jamadeo January 20, 2026 10:41

DOsinga approved these changes Jan 20, 2026

View reviewed changes

jamadeo reviewed Jan 20, 2026

View reviewed changes

jamadeo approved these changes Jan 21, 2026

View reviewed changes

lifeizhou-ap merged commit f26ffed into main Jan 21, 2026
20 checks passed

lifeizhou-ap deleted the lifei/fixed-accumulated-token-count branch January 21, 2026 16:27

github-actions bot mentioned this pull request Jan 22, 2026

chore(release): release version 1.21.0 (minor) #6634

Merged

fbalicchia pushed a commit to fbalicchia/goose that referenced this pull request Jan 23, 2026

Lifei/fixed accumulated token count (block#6587)

8b467e9

Signed-off-by: fbalicchia <fbalicchia@cuebiq.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Lifei/fixed accumulated token count #6587

Lifei/fixed accumulated token count #6587

Uh oh!

lifeizhou-ap commented Jan 20, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Jan 20, 2026

Uh oh!

Copilot AI Jan 20, 2026

Uh oh!

DOsinga Jan 20, 2026

Uh oh!

lifeizhou-ap Jan 20, 2026

Uh oh!

jamadeo left a comment

Uh oh!

lifeizhou-ap commented Jan 20, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

	if let Some(chunk_usage) = extract_usage_witqh_output_tokens(&tool_chunk) {
	if let Some(chunk_usage) = extract_usage_with_output_tokens(&tool_chunk) {

Lifei/fixed accumulated token count #6587

Lifei/fixed accumulated token count #6587

Uh oh!

Conversation

lifeizhou-ap commented Jan 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Type of Change

AI Assistance

Testing

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Jan 20, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 20, 2026

Choose a reason for hiding this comment

Uh oh!

DOsinga Jan 20, 2026

Choose a reason for hiding this comment

Uh oh!

lifeizhou-ap Jan 20, 2026

Choose a reason for hiding this comment

Uh oh!

jamadeo left a comment

Choose a reason for hiding this comment

Uh oh!

lifeizhou-ap commented Jan 20, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

lifeizhou-ap commented Jan 20, 2026 •

edited

Loading