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

[staging-supernova-dx-appf-io]

Summary

This PR captures streaming usage tokens and context management for Bedrock InvokeModel, ensuring consistent token accounting across both streaming and non-streaming paths.

Changes

Token Tracking Improvements

• Message attributes: Added stop_sequence and context_management to the Message class for richer response metadata
• Tokens class: Extended to track cache_creation_ephemeral_5m and cache_creation_ephemeral_1h in addition to existing cache token types
• StreamAccumulator: Added net_cache_tokens option to defer input_tokens netting, allowing proper calculation of effective input tokens after cache fields are updated by later stream chunks

Bedrock InvokeModel Streaming Enhancements

• streaming.rb:

  • Initialize StreamAccumulator with net_cache_tokens: true to handle the Bedrock pattern where message_start contains input_tokens but later message_delta chunks update cache fields
  • Capture cache-related tokens from stream chunks onto emitted events
  • Surface context_management.applied_edits from message_delta events

• chat.rb:

  • Extract and pass cache_creation usage fields to Message
  • Surface top-level context_management from response data

Input Tokens Semantics

Maintains consistent semantics: effective_input_tokens = max(input - cache_read - cache_creation, 0) across both streaming and non-streaming code paths.

Test Coverage

Added comprehensive specs for:

• Fully-cached request round trips where input_tokens reports the full cost but cached_tokens offset it
• Streaming message accumulation with cache token fields
• Context management propagation

References

• Fixes token accounting for Bedrock InvokeModel streaming responses
• Ensures cache tokens are properly accumulated and netted in final messages

---

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.