DUX-7485: Capture streaming usage tokens and context management for Bedrock InvokeModel

[staging-supernova-dx-appf-io]

Summary

Fix token usage capture inconsistency in Bedrock InvokeModel streaming protocol. The bug occurs because:

• message_start chunk sets input_tokens
• Later message_delta chunk can update cached_tokens and cache_creation_tokens
• Without synchronization, build_message_from_stream uses whichever chunk arrived last, causing inconsistent token accounting

Changes

• StreamAccumulator: Add net_cache_tokens option to defer input_tokens adjustment until to_message time. When enabled, input_tokens is netted against the final accumulated cached/cache_creation tokens to ensure consistent token accounting across streaming chunks.

• Message: Expose new fields stop_sequence and context_management from Bedrock InvokeModel responses.

• Tokens: Add cache creation TTL breakdown fields (cache_creation_ephemeral_5m, cache_creation_ephemeral_1h) to support granular cache telemetry.

• Bedrock InvokeModel chat: Extract and pass context_management and cache creation breakdown from response usage.

• Bedrock InvokeModel streaming: Initialize StreamAccumulator with net_cache_tokens: true to ensure token consistency.

Tests

• Add streaming usage token normalization coverage
• Verify distinct, non-zero usage on tool-call round trips
• Test context_management and stop_sequence capture in both streaming and non-streaming paths

---

Agent session: https://staging.supernova.dx.appf.io/coders/b96a7d9c-f609-466a-9a8e-e2020ad61ff8

[kermitapp]

Jira Issue: https://appfolio.atlassian.net/browse/DUX-7485

[staging-supernova-dx-appf-io]

🤖 Supernova Code Review — View trace

[staging-supernova-dx-appf-io]

Review summary

The PR cleanly wires up cache/ephemeral token extraction, stop_sequence, and context_management through the streaming path and aligns them with the non-streaming path. The test coverage is solid.

One bug: input_tokens can become inconsistent with cache tokens

The concrete issue is a split-responsibility design between build_message_start_chunk and build_message_delta_chunk:

• build_message_start_chunk pre-nets input_tokens (raw − cache_read − cache_creation) and passes only the netted value through the chunk. The raw total is discarded.
• build_message_delta_chunk (new in this PR) independently extracts cached_tokens and cache_creation_tokens from the message_delta event, which StreamAccumulator#count_tokens overwrites the accumulator's @cached_tokens/@cache_creation_tokens with.

If the message_delta event sends different cache counts than message_start — which this PR explicitly handles by reading them from both events — @input_tokens stays at its stale netted value while the cache totals change, violating:

input_tokens + cached_tokens + cache_creation_tokens == total_api_input_tokens

See inline comments on streaming.rb lines 183 and 281–284.

Recommended fix: store the raw API input_tokens in the accumulator (un-netted) and compute the net in to_message after all chunks have been accumulated:

# count_tokens: record raw, not netted
@raw_input_tokens = chunk.input_tokens + chunk.cached_tokens.to_i + chunk.cache_creation_tokens.to_i if chunk.input_tokens

# to_message: net down using final cache values
net_input = @raw_input_tokens ? [@raw_input_tokens - @cached_tokens.to_i - @cache_creation_tokens.to_i, 0].max : nil

This also requires build_message_start_chunk to pass the raw input_tokens value rather than the pre-netted one.

Note: StreamAccumulator#count_tokens (lines 165–172) is not in the diff so could not be commented inline — the coupling lives there.

---

Re-run this review

[staging-supernova-dx-appf-io]

input_tokens is netted down here using the message_start cache values, and the raw total is discarded. However, build_message_delta_chunk (lines 281–284) also extracts cached_tokens and cache_creation_tokens from the message_delta event, which the StreamAccumulator may overwrite @cached_tokens/@cache_creation_tokens with later. If the message_delta event ever reports different cache counts than message_start, the final message will violate the invariant input + cached + cache_creation == total_api_input_tokens.

Fix: pass the raw input_tokens through the chunk (e.g. as a new raw_input_tokens field) and defer the net-down to StreamAccumulator#to_message, using the final accumulated cache values at assembly time.

[staging-supernova-dx-appf-io]

These lines extract cached_tokens and cache_creation_tokens from message_delta, which StreamAccumulator#count_tokens will then overwrite the values set by message_start. This is the other side of the inconsistency described on line 183: @input_tokens is the pre-netted value from message_start, but @cached_tokens/@cache_creation_tokens can be independently updated here. If the two events ever disagree on cache counts, the accumulator will hold an inconsistent state where the stored input_tokens no longer matches total − cached − cache_creation.

[samuelboland]

Fix all issues

[staging-supernova-dx-appf-io]

@samuelboland

Fix all issues

Queued. The fix has two parts: (1) pass raw (un-netted) input_tokens from build_message_start_chunk in streaming.rb line 183, and (2) defer the netting computation to StreamAccumulator#to_message so it always uses the final accumulated cache values — eliminating the stale-@input_tokens inconsistency. The agent will also check other streaming providers before touching StreamAccumulator to make sure deferred netting doesn't break any provider that already passes pre-netted values.

---

🔍 View trace

[staging-supernova-dx-appf-io]

🤖 Supernova Code Review — View trace

[samuelboland]

Review please

[staging-supernova-dx-appf-io]

Reviewed the full diff across all 6 changed files. No actionable findings.

The design is sound:

• Deferred netting via net_cache_tokens: true — StreamAccumulator correctly holds @raw_input_tokens from message_start and applies max(input − cache_read − cache_creation, 0) only at to_message time, so a later message_delta that updates cache fields doesn't get missed.
• Ephemeral sub-tokens (cache_creation_ephemeral_5m, cache_creation_ephemeral_1h) are a breakdown of the existing cache_creation_input_tokens total, so excluding them from netted_input_tokens is correct — no double-counting.
• Non-streaming parity — chat.rb's input_tokens helper applies the same semantics and the new parse_completion_body fields (cache_creation_ephemeral_*, stop_sequence, context_management) are wired up consistently with Message and Tokens.
• Zero-value handling — Ruby's truthy semantics for 0 in count_tokens means a message_delta that explicitly sends cache_read_input_tokens: 0 would overwrite a prior non-zero value, which is the intended "last value wins" behaviour for Bedrock's cumulative update pattern.
• Test coverage is thorough for the new streaming paths: chunk-level extraction, full accumulator round-trips (fully-cached, partial-cache, cache-only-in-delta), and context_management propagation all have specs.

One minor observation (not blocking): there's no StreamAccumulator-level integration test that verifies cache_creation_ephemeral_5m/1h tokens survive the full add → to_message round trip (only the build_chunk unit test at line 530–545 covers the chunk level). The production path in count_cache_creation_ttl_tokens looks correct, but a matching accumulator-level spec alongside the one at line 614 would close the coverage gap and prevent regressions.

---

↩ Re-run this review

[staging-supernova-dx-appf-io]

@samuelboland

Review please

The fix is correct and the implementation is clean. Here's what's in the current state:

The original bug is resolved. build_message_start_chunk (line 183) now passes input_tok&.to_i — raw, un-netted. StreamAccumulator has a new net_cache_tokens: false constructor flag; the Bedrock InvokeModel path sets it true (streaming.rb line 24). netted_input_tokens (accumulator line 86-90) defers the subtraction to to_message time, using whatever final @cached_tokens/@cache_creation_tokens the accumulator holds after all chunks, whether those came from message_start or a later message_delta.

Test coverage is thorough. There are dedicated specs for:

• Raw input_tokens passing through un-netted at the chunk level (lines 493–511, 513–528)
• Correct end-to-end netting via net_cache_tokens: true (lines 614–643)
• Fully-cached round-trip yielding input_tokens == 0 (lines 645–664)
• The precise regression case: cache tokens arrive only on message_delta, not message_start (lines 666–690) — this is the right test to have

One minor observation: the end-to-end accumulation test at line 594–612 uses StreamAccumulator.new without net_cache_tokens: true. That's intentional since it's simulating a non-Bedrock-InvokeModel provider, but a comment noting that would help future readers understand why it diverges from the Bedrock path tests.

Nothing blocking. The PR looks good to merge.

[staging-supernova-dx-appf-io]

🤖 Supernova Code Review — View trace

[staging-supernova-dx-appf-io]

Review Summary

The token-accumulation and netting logic looks correct overall. A few notes:

Dropped inline finding (line not in diff) — lib/ruby_llm/message.rb

The new cache_creation_ephemeral_5m and cache_creation_ephemeral_1h fields on Tokens are surfaced in to_h and accessible on Tokens, but unlike every other token sub-field (cached_tokens, cache_creation_tokens, thinking_tokens, etc.) they have no delegate accessor on Message. Callers must use message.tokens&.cache_creation_ephemeral_5m rather than the consistent message.cache_creation_ephemeral_5m_tokens pattern. Consider adding:

def cache_creation_ephemeral_5m_tokens
  tokens&.cache_creation_ephemeral_5m
end

def cache_creation_ephemeral_1h_tokens
  tokens&.cache_creation_ephemeral_1h
end

Deferred-netting design (net_cache_tokens: true)

The approach is sound: message_start carries the raw input_tokens and initial cache snapshot, while a later message_delta may refine the cache counts; netted_input_tokens correctly waits until to_message to apply max(input - cached - cache_creation, 0). One edge case to be aware of: the guard if chunk.cached_tokens in count_tokens will update @cached_tokens to 0 if Bedrock ever sends an explicit "cache_read_input_tokens": 0 in message_delta (since 0 is truthy in Ruby), potentially overriding a non-zero value from message_start. The current test suite only covers cases where the delta either omits the field or repeats the same value. If Bedrock can send a contradicting explicit zero here, a test covering that scenario would be valuable.

No blocking issues found.

---

↩ Re-run this review

[staging-supernova-dx-appf-io]

nit: cached_tokens: and cache_creation_tokens: are passed as raw JSON values here, but input_tokens: gets an explicit .to_i call. Since JSON parsing already yields integers, this isn't a crash risk, but the inconsistency is worth noting — if the pattern ever changes or a non-numeric value slips through, netted_input_tokens's .to_i fallback in stream_accumulator.rb masks the mismatch rather than raising early.

Consider applying .to_i to all three token values for consistency with the input_tok&.to_i pattern directly above.