You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[ML-70509] Fix streamed and cached token accounting in ChatDatabricks - #746
When LangGraph streams ChatDatabricks messages, cumulative usage attached to each content chunk is added again with the final usage chunk. Three snapshots of 2,697 input tokens therefore become 10,788. Claude cache counters also inflate input_tokens because the gateway's OpenAI-compatible prompt_tokens already includes cached inputs. GPT cache-write counts are discarded in both API converters.
This change emits the final streaming usage snapshot once in synchronous and asynchronous calls, preserves the gateway's inclusive input counts, and maps nonzero cache writes to LangChain's cache_creation detail for Chat Completions and Responses. stream_usage=True requests include_usage for Chat Completions while preserving explicit stream_options overrides, including include_usage=False. The streaming documentation now aggregates the stream before reading usage.
The converters follow the OpenAI-compatible gateway contract. Native Anthropic input_tokens is exclusive of cache reads/writes and is a different API schema; these converters do not infer a provider or change counts using model-name or arithmetic heuristics.
Validation:
45 transport-level regression cases cover JSON/SSE parsing, sync/async invoke and streaming, usage-only trailers, cumulative snapshots, explicit opt-outs, absent/zero/nonzero cache writes, and LangGraph messages+values aggregation. The original 2,697→10,788 symptom reproduces before the source fix.
Full LangChain unit suite: 345 passed, 2 skipped, 1 unrelated timeout. test_vector_search_client_model_serving_environment retries DNS for its testdogfod.com fixture with the installed Databricks SDK. The same test times out against unmodified main; all chat model and usage tests pass.
Ruff check and format pass. ty check exits successfully with two existing unused-ignore warnings.
Bounded live DEFAULT verification against the same Claude request: three provider usage snapshots produce one usage chunk; input 3,011, output 5, total 3,016 and cache_read 3,002 exactly match the final provider payload. A preceding live request confirmed cache creation 3,002 is already included in prompt_tokens 3,011.
This draft overlaps the streaming fix in open #638 and adds the remaining accounting changes and coverage for ML-70509. Open #388 concerns exposing cache counters in raw llm_output metadata and remains separate. Merged #453 only updated zero-valued SDK fixtures; it did not preserve cache-write usage. Publishing a new databricks-langchain version and updating Agent Bricks' 0.21.0 pins will be needed to deliver this fix to generated projects.
Optional private-runner integration tests have not been triggered.
CI framework limitation: the only failing test cases are the six memory-backed test_history_orders_parallel_pending_messages variants, with pending-write order ["first", "last"] instead of ["last", "first"]. The identical six cases also fail on base commit a8d55f4 in main run 38662891535. They are pre-existing; no skips or forced passes were applied.
Combined private cache validation: https://github.com/databricks-eng/ai-oss-integration-tests-runner/actions/runs/38698567494 passed all 20 live caching tests (OpenAI 8, LangChain 4, template adapters 8) at e231ca2. This source combines 11 PR heads; it is not an isolated run of this branch. Models cover Claude Sonnet 4.5/4.6 in the template adapters and Gateway/FMAPI routes in the connectors. The Apps/tool matrix runs separately.
When combining #727 and #750, compose the Gemini signature and cache normalizers in each shared OpenAI response hook. Keep #746 usage accounting with the cache regressions in the LangChain merge. The test-only combined commit records those resolutions; individual PRs retain independent boundaries.
Bug-bash integration validation
The test-only combined branch at be2fa8ee8d6d3b3e011ecfa285e6455b90f9433a contains this exact PR head and eleven other reviewed heads. Its Agent Bricks unit suite passed 2,154 tests, 24 skipped. This is combined-source validation, not an isolated run of this PR; the branch must not be merged into main.
Private run 38698567494 passed all 20 live Claude cache tests at e231ca2 (OpenAI 8, LangGraph 4, generated adapters 8). Run 38701390693 at 53f9db8 passed OpenAI 8, LangGraph 4 and 3 adapter cases; the remaining 5 adapter cases returned HTTP403 because the shared test-workspace Gateway budget reached its $10 limit. Production source is identical across those runs and the final combined source; later changes affect only matrix fixtures, cleanup, regression tests and documentation. The rerun is not all green, and no budgets were changed.
The corrected Apps test passed 20/20 matrix rows, two grant snapshots, transitive-permission controls and repeat-deploy idempotency. Its evidence confirms clean source and wheel/source matching. That live run predates the final cleanup correction: the harness recorded twelve successful cleanup operations but omitted the implicit Runtime Store/role for direct-authoring Apps. #752 now includes ownership-checked cleanup for every App and safe retry/evidence handling, covered by focused regressions and the final unit suite. No new full model matrix was run for those cleanup-only changes.
Independent read-only run 38708044949 confirmed all 13 exact prior-run resources absent, including three consecutive typed404 observations for the omitted Runtime Store, dedicated database and role. One separate current-window database/role pair remains present and unattributed: SCIM did not recover its App identity and one bounded audit SELECT failed. That pair was left intact; full cleanup of the current matrix run is not established. No resources were deleted using name/time correlation alone.
The Apps matrix covers LangGraph, CLI/direct authoring, dev/deployed runtime and five tool kinds. Separate cache tests cover both LangGraph and OpenAI Agents. These checks do not establish fresh-account provisioning rights or freshly consented deployed Genie OBO acceptance.
All public checks finished. #694/#727 have 54 green; the other ten PRs have 53 green and the same framework job red. Its six memory-backed parallel-history ordering failures also occur on current main. No forced skips or weakened grant assertions were applied.
jamesbxwu
changed the title
Fix streamed and cached token accounting in ChatDatabricks
[ML-70509] Fix streamed and cached token accounting in ChatDatabricks
Oct 15, 2026
Closing as consolidated into #727. The sync/async streaming deduplication also overlaps the earlier #638; #727 already included that behavior and the inclusive Claude input-total correction.
The remaining behavior from this draft is now in #727 head 3963189: cache_write_tokens extraction for Chat Completions and Responses, default streaming usage requests with explicit opt-outs preserved, and HTTP-mocked transport/LangGraph aggregation regressions. Partial native/unified cache counters merge independently in LangChain and DatabricksOpenAI, including explicit zeros.
Current focused validation: 177 LangChain tests passed, 2 skipped; 15 OpenAI normalization tests passed; Ruff, formatting, focused types and secret checks passed. No new live/model calls were run. ML-70509 remains pending merge/release; closing this duplicate does not mean the fix is released. No PR was merged into main.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #434 and #446.
When LangGraph streams ChatDatabricks messages, cumulative usage attached to each content chunk is added again with the final usage chunk. Three snapshots of 2,697 input tokens therefore become 10,788. Claude cache counters also inflate input_tokens because the gateway's OpenAI-compatible prompt_tokens already includes cached inputs. GPT cache-write counts are discarded in both API converters.
This change emits the final streaming usage snapshot once in synchronous and asynchronous calls, preserves the gateway's inclusive input counts, and maps nonzero cache writes to LangChain's cache_creation detail for Chat Completions and Responses. stream_usage=True requests include_usage for Chat Completions while preserving explicit stream_options overrides, including include_usage=False. The streaming documentation now aggregates the stream before reading usage.
The converters follow the OpenAI-compatible gateway contract. Native Anthropic input_tokens is exclusive of cache reads/writes and is a different API schema; these converters do not infer a provider or change counts using model-name or arithmetic heuristics.
Validation:
This draft overlaps the streaming fix in open #638 and adds the remaining accounting changes and coverage for ML-70509. Open #388 concerns exposing cache counters in raw llm_output metadata and remains separate. Merged #453 only updated zero-valued SDK fixtures; it did not preserve cache-write usage. Publishing a new databricks-langchain version and updating Agent Bricks' 0.21.0 pins will be needed to deliver this fix to generated projects.
Optional private-runner integration tests have not been triggered.
CI framework limitation: the only failing test cases are the six memory-backed
test_history_orders_parallel_pending_messagesvariants, with pending-write order["first", "last"]instead of["last", "first"]. The identical six cases also fail on base commit a8d55f4 in main run 38662891535. They are pre-existing; no skips or forced passes were applied.Combined private cache validation: https://github.com/databricks-eng/ai-oss-integration-tests-runner/actions/runs/38698567494 passed all 20 live caching tests (OpenAI 8, LangChain 4, template adapters 8) at e231ca2. This source combines 11 PR heads; it is not an isolated run of this branch. Models cover Claude Sonnet 4.5/4.6 in the template adapters and Gateway/FMAPI routes in the connectors. The Apps/tool matrix runs separately.
When combining #727 and #750, compose the Gemini signature and cache normalizers in each shared OpenAI response hook. Keep #746 usage accounting with the cache regressions in the LangChain merge. The test-only combined commit records those resolutions; individual PRs retain independent boundaries.
Bug-bash integration validation
The test-only combined branch at
be2fa8ee8d6d3b3e011ecfa285e6455b90f9433acontains this exact PR head and eleven other reviewed heads. Its Agent Bricks unit suite passed 2,154 tests, 24 skipped. This is combined-source validation, not an isolated run of this PR; the branch must not be merged into main.Private run 38698567494 passed all 20 live Claude cache tests at
e231ca2(OpenAI 8, LangGraph 4, generated adapters 8). Run 38701390693 at53f9db8passed OpenAI 8, LangGraph 4 and 3 adapter cases; the remaining 5 adapter cases returned HTTP403 because the shared test-workspace Gateway budget reached its $10 limit. Production source is identical across those runs and the final combined source; later changes affect only matrix fixtures, cleanup, regression tests and documentation. The rerun is not all green, and no budgets were changed.The corrected Apps test passed 20/20 matrix rows, two grant snapshots, transitive-permission controls and repeat-deploy idempotency. Its evidence confirms clean source and wheel/source matching. That live run predates the final cleanup correction: the harness recorded twelve successful cleanup operations but omitted the implicit Runtime Store/role for direct-authoring Apps. #752 now includes ownership-checked cleanup for every App and safe retry/evidence handling, covered by focused regressions and the final unit suite. No new full model matrix was run for those cleanup-only changes.
Independent read-only run 38708044949 confirmed all 13 exact prior-run resources absent, including three consecutive typed404 observations for the omitted Runtime Store, dedicated database and role. One separate current-window database/role pair remains present and unattributed: SCIM did not recover its App identity and one bounded audit SELECT failed. That pair was left intact; full cleanup of the current matrix run is not established. No resources were deleted using name/time correlation alone.
The Apps matrix covers LangGraph, CLI/direct authoring, dev/deployed runtime and five tool kinds. Separate cache tests cover both LangGraph and OpenAI Agents. These checks do not establish fresh-account provisioning rights or freshly consented deployed Genie OBO acceptance.
All public checks finished. #694/#727 have 54 green; the other ten PRs have 53 green and the same framework job red. Its six memory-backed parallel-history ordering failures also occur on current main. No forced skips or weakened grant assertions were applied.