Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
staging-supernova-dx-appf-io wants to merge 3 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
16c1731
Capture streaming usage tokens and context management for Bedrock Inv…
appf-supernova Jul 6, 2026
7bdc118
Defer input_tokens netting to StreamAccumulator for Bedrock InvokeModel
appf-supernova Jul 7, 2026
6213d75
Add test coverage for streaming usage token normalization
appf-supernova Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 9 additions & 3 deletions lib/ruby_llm/message.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ class Message
:CONTENT_FILTERED_FINISH_REASONS

attr_reader :role, :model_id, :tool_calls, :tool_call_id, :raw, :thinking, :tokens, :citations,
:finish_reason
:finish_reason, :stop_sequence, :context_management
attr_writer :content

def initialize(options = {})
Expand All @@ -29,12 +29,16 @@ def initialize(options = {})
output: options[:output_tokens],
cached: options[:cached_tokens],
cache_creation: options[:cache_creation_tokens],
thinking: options[:thinking_tokens]
thinking: options[:thinking_tokens],
cache_creation_ephemeral_5m: options[:cache_creation_ephemeral_5m_tokens],
cache_creation_ephemeral_1h: options[:cache_creation_ephemeral_1h_tokens]
)
@raw = options[:raw]
@thinking = options[:thinking]
@citations = Array(options[:citations])
@finish_reason = options[:finish_reason]
@stop_sequence = options[:stop_sequence]
@context_management = options[:context_management]

ensure_valid_role
end
Expand Down Expand Up @@ -113,7 +117,9 @@ def to_h
thinking: thinking&.text,
thinking_signature: thinking&.signature,
citations: citations.empty? ? nil : citations.map(&:to_h),
finish_reason: finish_reason
finish_reason: finish_reason,
stop_sequence: stop_sequence,
context_management: context_management
}.merge(tokens ? tokens.to_h : {}).compact
end

Expand Down
5 changes: 5 additions & 0 deletions lib/ruby_llm/protocols/bedrock_invoke_model/chat.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,7 @@ def parse_completion_body(data, raw:)

content_blocks = data['content'] || []
usage = data['usage'] || {}
cache_creation = usage['cache_creation'] || {}

Message.new(
role: :assistant,
Expand All @@ -115,7 +116,11 @@ def parse_completion_body(data, raw:)
output_tokens: usage['output_tokens'],
cached_tokens: usage['cache_read_input_tokens'],
cache_creation_tokens: usage['cache_creation_input_tokens'],
cache_creation_ephemeral_5m_tokens: cache_creation['ephemeral_5m_input_tokens'],
cache_creation_ephemeral_1h_tokens: cache_creation['ephemeral_1h_input_tokens'],
finish_reason: data['stop_reason'],
stop_sequence: data['stop_sequence'],
context_management: data['context_management'],
model_id: data['model'],
raw: raw
)
Expand Down
30 changes: 27 additions & 3 deletions lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ def stream_url
end

def stream_response(payload, additional_headers = {}, &block)
accumulator = StreamAccumulator.new
accumulator = StreamAccumulator.new(net_cache_tokens: true)
decoder = event_stream_decoder
thinking_state = {}
body = JSON.generate(payload)
Expand All @@ -48,9 +48,19 @@ def stream_response(payload, additional_headers = {}, &block)

message = accumulator.to_message(response)
RubyLLM.logger.debug { "Stream completed: #{message.content}" }
log_context_management(message)
message
end

def log_context_management(message)
applied_edits = message.context_management && message.context_management['applied_edits']
return unless applied_edits && !applied_edits.empty?

RubyLLM.logger.debug do
"Bedrock InvokeModel context_management applied_edits: #{applied_edits.inspect}"
end
end

def event_stream_decoder
require 'aws-eventstream'
Aws::EventStream::Decoder.new
Expand Down Expand Up @@ -161,13 +171,20 @@ def build_chunk(event, thinking_state = {})
def build_message_start_chunk(event)
message = event['message'] || {}
usage = message['usage'] || {}
cache_read = usage['cache_read_input_tokens']
cache_creation = usage['cache_creation_input_tokens']
input_tok = usage['input_tokens']
cache_creation_detail = usage['cache_creation'] || {}

Chunk.new(
role: :assistant,
content: nil,
model_id: message['model'] || @model&.id,
input_tokens: input_tok ? [input_tok.to_i, 0].max : nil
input_tokens: input_tok&.to_i,
cached_tokens: cache_read,
cache_creation_tokens: cache_creation,
cache_creation_ephemeral_5m_tokens: cache_creation_detail['ephemeral_5m_input_tokens'],
cache_creation_ephemeral_1h_tokens: cache_creation_detail['ephemeral_1h_input_tokens']

@staging-supernova-dx-appf-io staging-supernova-dx-appf-io Bot Jul 8, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

)
end

Expand Down Expand Up @@ -254,13 +271,20 @@ def build_content_block_stop_chunk(event, thinking_state)
def build_message_delta_chunk(event)
delta = event['delta'] || {}
usage = event['usage'] || {}
cache_creation_detail = usage['cache_creation'] || {}

Chunk.new(
role: :assistant,
content: nil,
model_id: @model&.id,
output_tokens: usage['output_tokens'],
finish_reason: delta['stop_reason']
cached_tokens: usage['cache_read_input_tokens'],
cache_creation_tokens: usage['cache_creation_input_tokens'],
cache_creation_ephemeral_5m_tokens: cache_creation_detail['ephemeral_5m_input_tokens'],
cache_creation_ephemeral_1h_tokens: cache_creation_detail['ephemeral_1h_input_tokens'],
Comment on lines +281 to +284

@staging-supernova-dx-appf-io staging-supernova-dx-appf-io Bot Jul 6, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

finish_reason: delta['stop_reason'],
stop_sequence: delta['stop_sequence'],
context_management: delta['context_management']
)
end

Expand Down
42 changes: 37 additions & 5 deletions lib/ruby_llm/stream_accumulator.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,19 +8,31 @@ module RubyLLM
class StreamAccumulator
attr_reader :content, :model_id, :tool_calls

def initialize
# net_cache_tokens: when true, input_tokens is netted against the final accumulated
# cached_tokens/cache_creation_tokens at to_message time instead of being taken as-is
# from whichever chunk last set it. Bedrock InvokeModel's message_start chunk carries
# input_tokens while a later message_delta chunk can still update the cache fields, so
# netting must be deferred until all chunks are in. Other providers already net
# input_tokens against cache fields within the same usage snapshot before building the
# chunk, so deferred netting stays off for them to avoid double-subtracting.
def initialize(net_cache_tokens: false)
@net_cache_tokens = net_cache_tokens
@content = +''
@citations = []
@thinking_text = +''
@thinking_signature = nil
@thinking_blocks = []
@tool_calls = {}
@input_tokens = nil
@raw_input_tokens = nil
@output_tokens = nil
@cached_tokens = nil
@cache_creation_tokens = nil
@cache_creation_ephemeral_5m_tokens = nil
@cache_creation_ephemeral_1h_tokens = nil
@thinking_tokens = nil
@finish_reason = nil
@stop_sequence = nil
@context_management = nil
@inside_think_tag = false
@pending_think_tag = +''
@latest_tool_call_id = nil
Expand All @@ -35,6 +47,8 @@ def add(chunk)
accumulate_citations(chunk.citations)
append_thinking_from_chunk(chunk)
@finish_reason = chunk.finish_reason if chunk.finish_reason
@stop_sequence = chunk.stop_sequence if chunk.stop_sequence
@context_management = chunk.context_management if chunk.context_management
count_tokens chunk
RubyLLM.logger.debug { inspect } if RubyLLM.config.log_stream_debug
end
Expand All @@ -50,13 +64,17 @@ def to_message(response)
blocks: @thinking_blocks.empty? ? nil : @thinking_blocks
),
tokens: Tokens.build(
input: @input_tokens,
input: netted_input_tokens,
output: @output_tokens,
cached: @cached_tokens,
cache_creation: @cache_creation_tokens,
thinking: @thinking_tokens
thinking: @thinking_tokens,
cache_creation_ephemeral_5m: @cache_creation_ephemeral_5m_tokens,
cache_creation_ephemeral_1h: @cache_creation_ephemeral_1h_tokens
),
finish_reason: @finish_reason,
stop_sequence: @stop_sequence,
context_management: @context_management,
model_id: model_id,
tool_calls: tool_calls_from_stream,
raw: response
Expand All @@ -65,6 +83,12 @@ def to_message(response)

private

def netted_input_tokens
return @raw_input_tokens unless @net_cache_tokens && @raw_input_tokens

[@raw_input_tokens - @cached_tokens.to_i - @cache_creation_tokens.to_i, 0].max
end

# Providers like Perplexity repeat the full citation list on every chunk.
def accumulate_citations(new_citations)
new_citations.each do |citation|
Expand Down Expand Up @@ -153,11 +177,19 @@ def find_tool_call(stream_key)
end

def count_tokens(chunk)
@input_tokens = chunk.input_tokens if chunk.input_tokens
@raw_input_tokens = chunk.input_tokens if chunk.input_tokens
@output_tokens = chunk.output_tokens if chunk.output_tokens
@cached_tokens = chunk.cached_tokens if chunk.cached_tokens
@cache_creation_tokens = chunk.cache_creation_tokens if chunk.cache_creation_tokens
@thinking_tokens = chunk.thinking_tokens if chunk.thinking_tokens
count_cache_creation_ttl_tokens(chunk)
end

def count_cache_creation_ttl_tokens(chunk)
ephemeral_5m = chunk.tokens&.cache_creation_ephemeral_5m
ephemeral_1h = chunk.tokens&.cache_creation_ephemeral_1h
@cache_creation_ephemeral_5m_tokens = ephemeral_5m if ephemeral_5m
@cache_creation_ephemeral_1h_tokens = ephemeral_1h if ephemeral_1h
end

def handle_chunk_content(chunk)
Expand Down
26 changes: 20 additions & 6 deletions lib/ruby_llm/tokens.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,35 +3,49 @@
module RubyLLM
# Represents token usage for a response.
class Tokens
attr_reader :input, :output, :cached, :cache_creation, :thinking
attr_reader :input, :output, :cached, :cache_creation, :thinking,
:cache_creation_ephemeral_5m, :cache_creation_ephemeral_1h

def initialize(input: nil, output: nil, cached: nil, cache_creation: nil, thinking: nil)
# rubocop:disable Metrics/ParameterLists
def initialize(input: nil, output: nil, cached: nil, cache_creation: nil, thinking: nil,
cache_creation_ephemeral_5m: nil, cache_creation_ephemeral_1h: nil)
@input = input
@output = output
@cached = cached
@cache_creation = cache_creation
@thinking = thinking
@cache_creation_ephemeral_5m = cache_creation_ephemeral_5m
@cache_creation_ephemeral_1h = cache_creation_ephemeral_1h
end
# rubocop:enable Metrics/ParameterLists

def self.build(input: nil, output: nil, cached: nil, cache_creation: nil, thinking: nil)
return nil if [input, output, cached, cache_creation, thinking].all?(&:nil?)
# rubocop:disable Metrics/ParameterLists
def self.build(input: nil, output: nil, cached: nil, cache_creation: nil, thinking: nil,
cache_creation_ephemeral_5m: nil, cache_creation_ephemeral_1h: nil)
return nil if [input, output, cached, cache_creation, thinking,
cache_creation_ephemeral_5m, cache_creation_ephemeral_1h].all?(&:nil?)

new(
input: input,
output: output,
cached: cached,
cache_creation: cache_creation,
thinking: thinking
thinking: thinking,
cache_creation_ephemeral_5m: cache_creation_ephemeral_5m,
cache_creation_ephemeral_1h: cache_creation_ephemeral_1h
)
end
# rubocop:enable Metrics/ParameterLists

def to_h
{
input_tokens: input,
output_tokens: output,
cached_tokens: cached,
cache_creation_tokens: cache_creation,
thinking_tokens: thinking
thinking_tokens: thinking,
cache_creation_ephemeral_5m_tokens: cache_creation_ephemeral_5m,
cache_creation_ephemeral_1h_tokens: cache_creation_ephemeral_1h
}.compact
end

Expand Down
Loading