fix(ruby_llm): read token counts from RubyLLM 2.x turns

CHANGELOG.md

@@ -1,5 +1,21 @@
 # Changelog
 
+## [0.3.1] - 2026-09-20
+
+### Fixed
+
+- The RubyLLM adapter reads token counts from a RubyLLM 2.x turn. 2.0 moved
+  the per-message readers (`input_tokens`, `output_tokens`, `thinking_tokens`)
+  to `message.tokens` and puts the round's `Tokens` on the `chat.ruby_llm`
+  payload; the adapter only knew the 1.x readers, so every 2.x trace reported
+  0 tokens while its spans, timings and errors were correct. The payload's
+  `tokens` is read first, then the round's messages with either generation's
+  readers, so 1.x traces are unchanged.
+
+The core gem is republished at 0.3.1 with no changes, so the two gems stay on
+one version.
+
+
 ## [0.3.0] - 2026-09-12
 
 ### Added

adapters/ruby_llm/lib/activeagents/telemetry/ruby_llm.rb

@@ -43,6 +43,9 @@ module RubyLLM
 
       DEFAULT_AGENT = { name: "RubyLLM::Chat", action: "chat" }.freeze
 
+      # Span token keys and the reader that answers each on a RubyLLM 2.x +Tokens+.
+      TOKEN_READERS = { "input" => :input, "output" => :output, "thinking" => :thinking }.freeze
+
       State = Struct.new(:depth, :started_at, :tool_spans, :rounds, :tokens, :chat_key, :agent)
 
       class << self
@@ -306,22 +309,39 @@ def truncate_captured(text)
           text.to_s[0, CONTENT_LIMIT]
         end
 
+        # Tokens for one round. RubyLLM 2.x puts the round's +Tokens+ on the
+        # payload and moves the per-message readers to +message.tokens+; 1.x
+        # has +input_tokens+ and friends on each assistant message and nothing
+        # on the payload. The payload is read first, then the messages the
+        # round added, with either generation's readers.
         def token_totals(payload)
+          tokens = payload_tokens(payload[:tokens]) || message_tokens(payload)
+          tokens["total"] = tokens.values.sum
+          tokens
+        end
+
+        def payload_tokens(tokens)
+          return unless TOKEN_READERS.values.all? { |reader| tokens.respond_to?(reader) }
+
+          TOKEN_READERS.to_h { |key, reader| [ key, tokens.public_send(reader).to_i ] }
+        end
+
+        def message_tokens(payload)
           initial_count = Array(payload[:input_messages]).size
           new_messages = Array(payload[:messages_after])[initial_count..] || []
           assistant_messages = new_messages.select { |message| message.respond_to?(:role) && message.role.to_s == "assistant" }
 
-          tokens = {
-            "input" => sum_tokens(assistant_messages, :input_tokens),
-            "output" => sum_tokens(assistant_messages, :output_tokens),
-            "thinking" => sum_tokens(assistant_messages, :thinking_tokens)
-          }
-          tokens["total"] = tokens.values.sum
-          tokens
+          TOKEN_READERS.to_h { |key, reader| [ key, assistant_messages.sum { |message| message_token_count(message, key, reader) } ] }
         end
 
-        def sum_tokens(messages, method_name)
-          messages.sum { |message| message.respond_to?(method_name) ? message.public_send(method_name).to_i : 0 }
+        # A 1.x message answers +input_tokens+; a 2.x message answers
+        # +tokens.input+.
+        def message_token_count(message, key, reader)
+          legacy_reader = :"#{key}_tokens"
+          return message.public_send(legacy_reader).to_i if message.respond_to?(legacy_reader)
+
+          tokens = message.tokens if message.respond_to?(:tokens)
+          tokens.respond_to?(reader) ? tokens.public_send(reader).to_i : 0
         end
       end
 

adapters/ruby_llm/lib/activeagents/telemetry/ruby_llm/version.rb

@@ -3,7 +3,7 @@
 module ActiveAgents
   module Telemetry
     module RubyLLM
-      VERSION = "0.3.0"
+      VERSION = "0.3.1"
     end
   end
 end

adapters/ruby_llm/test/test_ruby_llm.rb

@@ -144,6 +144,49 @@ def test_inherits_destination_from_the_shared_configuration
     ActiveAgents::Telemetry.reset!
   end
 
+  # RubyLLM 2.x: the round's Tokens ride the payload; messages answer #tokens, not #input_tokens.
+  Tokens2 = Struct.new(:input, :output, :thinking, keyword_init: true)
+  Msg2 = Struct.new(:role, :tokens)
+
+  def test_reads_a_2x_rounds_tokens_from_the_payload
+    reply = Msg2.new("assistant", Tokens2.new(input: 10, output: 5, thinking: 2))
+    payload = chat_payload(input_messages: [ Msg2.new("user", nil) ], tokens: Tokens2.new(input: 10, output: 5, thinking: 2))
+
+    instrument("chat.ruby_llm", payload) { |p| p[:messages_after] = [ *payload[:input_messages], reply ]; p[:tool_call] = false }
+
+    assert_equal({ "input" => 10, "output" => 5, "thinking" => 2, "total" => 17 }, spans_of(traces.last, "llm").first["tokens"])
+  end
+
+  def test_sums_2x_payload_tokens_across_sibling_rounds
+    chat = Object.new
+    first = chat_payload(chat: chat, tokens: Tokens2.new(input: 10, output: 5, thinking: 0))
+    instrument("chat.ruby_llm", first) { |p| p[:messages_after] = [ *first[:input_messages], Msg2.new("assistant", first[:tokens]) ]; p[:tool_call] = true }
+    second = chat_payload(chat: chat, tokens: Tokens2.new(input: 20, output: 7, thinking: 0))
+    instrument("chat.ruby_llm", second) { |p| p[:messages_after] = [ *second[:input_messages], Msg2.new("assistant", second[:tokens]) ]; p[:tool_call] = false }
+
+    assert_equal 1, traces.size
+    llm_span = spans_of(traces.last, "llm").first
+    assert_equal 2, llm_span["attributes"]["llm.rounds"]
+    assert_equal({ "input" => 30, "output" => 12, "thinking" => 0, "total" => 42 }, llm_span["tokens"])
+  end
+
+  def test_reads_2x_message_tokens_when_the_payload_carries_none
+    reply = Msg2.new("assistant", Tokens2.new(input: 10, output: 5, thinking: 0))
+    payload = chat_payload(input_messages: [ Msg2.new("user", nil) ])
+
+    instrument("chat.ruby_llm", payload) { |p| p[:messages_after] = [ *payload[:input_messages], reply ]; p[:tool_call] = false }
+
+    assert_equal({ "input" => 10, "output" => 5, "thinking" => 0, "total" => 15 }, spans_of(traces.last, "llm").first["tokens"])
+  end
+
+  def test_an_empty_2x_tokens_object_reports_zero_rather_than_raising
+    payload = chat_payload(tokens: Tokens2.new(input: nil, output: nil, thinking: nil))
+
+    instrument("chat.ruby_llm", payload) { |p| p[:messages_after] = payload[:input_messages]; p[:tool_call] = false }
+
+    assert_equal({ "input" => 0, "output" => 0, "thinking" => 0, "total" => 0 }, spans_of(traces.last, "llm").first["tokens"])
+  end
+
   ContentMsg = Struct.new(:role, :content, :input_tokens, :output_tokens, :thinking_tokens)
 
   def test_does_not_capture_content_by_default

docs/branches/ruby-llm-2x-tokens.md

@@ -0,0 +1,6 @@
+# Branch
+
+- Branch: `claude/ruby-llm-2x-tokens`.
+- Base: `main` at `f401a02f30c18196a109b59fe071da04d539dbcc`.
+- Scope: RubyLLM 2.x token counts in the adapter; both gems bumped to 0.3.1.
+- Release: tag `v0.3.1` after the merge; the release workflow publishes both gems.

docs/issues/ruby-llm-2x-tokens.md

@@ -0,0 +1,11 @@
+# RubyLLM 2.x turns report zero tokens
+
+RubyLLM 2.0 moved the per-message token readers (`input_tokens`,
+`output_tokens`, `thinking_tokens`) to `message.tokens` and put the round's
+`Tokens` on the `chat.ruby_llm` payload. The adapter only knew the 1.x readers,
+so every 2.x turn reported `0` input and output tokens while its spans, timings
+and errors were correct. A dashboard fed by a RubyLLM 2.x app showed no token
+usage and no cost.
+
+The adapter now reads the payload's `tokens` first, then the round's messages
+with either generation's readers. 1.x traces are unchanged.

docs/milestones/ruby-llm-2x-tokens.md

@@ -0,0 +1,7 @@
+# Milestone: RubyLLM 2.x token counts
+
+- [x] Read the round's `Tokens` from the 2.x payload, then the messages with either generation's readers.
+- [x] Test the 2.x payload, sibling rounds, the message-only fallback and an empty `Tokens`.
+- [x] Prove the fix in a booted Rails app on RubyLLM 1.16 and 2.0, from the packaged gems.
+- [x] Bump both gems to 0.3.1 with a changelog entry, so the merge is releasable.
+- [ ] Merge the PR, then tag `v0.3.1` to publish both gems through the release workflow.

docs/pull-requests/ruby-llm-2x-tokens.md

@@ -0,0 +1,16 @@
+# PR: RubyLLM 2.x token counts
+
+Branch `claude/ruby-llm-2x-tokens` makes the adapter read token counts from a
+RubyLLM 2.x turn. 2.0 moved the per-message readers to `message.tokens` and
+puts the round's `Tokens` on the `chat.ruby_llm` payload; the adapter only knew
+the 1.x readers, so every 2.x trace reported 0 tokens. The payload's `tokens`
+is read first, then the round's messages with either generation's readers, so
+1.x traces are unchanged.
+
+Core and adapter suites pass: 63 tests, 177 assertions, zero failures. The fix
+was also run in a Rails 8.1 app mounting the dashboard, on RubyLLM 1.16.0 and
+2.0.0, from the built gems: with the released adapter every 2.x trace stored
+`0`/`0` tokens; with 0.3.1 a plain chat stores 11/7, and a two-round tool turn
+22/14. Versions and the changelog are bumped to 0.3.1 on the branch (the core
+gem is republished unchanged so the two stay on one version); tagging `v0.3.1`
+after the merge publishes both gems.

lib/activeagents/telemetry/version.rb

@@ -2,6 +2,6 @@
 
 module ActiveAgents
   module Telemetry
-    VERSION = "0.3.0"
+    VERSION = "0.3.1"
   end
 end