Skip to content

[ML-70509] Fix streamed and cached token accounting in ChatDatabricks - #746

Closed
jamesbxwu wants to merge 1 commit into
mainfrom
codex/ml-70509-streaming-usage
Closed

jamesbxwu wants to merge 1 commit into
mainfrom
codex/ml-70509-streaming-usage

Conversation

@jamesbxwu

@jamesbxwu jamesbxwu commented Oct 15, 2026 •

Copy link
Copy Markdown
Contributor

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:

  • 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 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
@jamesbxwu

Copy link
Copy Markdown
Contributor Author

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.

@jamesbxwu jamesbxwu closed this Oct 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[ChatDatabricks] Token usage counts are inflated when using streaming

1 participant