token-lens 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: fc58978bdc3f4759309ded99d65da22be12a01287ea79acdb43ab873e7cb8771
+  data.tar.gz: 34c54363d13e72ed669e3f27bbf1fc8af2d3665e958d977baabfaac9ab80f849
+SHA512:
+  metadata.gz: 3f3c35dedc0cd0efd6e12b8d136aab7f4328ba4619f912f07c94727e3bd6c4c2bb4184d07a9e3bbb117dd0eaeec28edd27c5f7100f140a79b12043773a066881
+  data.tar.gz: d0767d26c77b4ba6dfec6f3a40b70c7d1dde3b3a58094fa8707b4bf57bc860612ed1af91c161b31fc3f3faa78f3db25ed6f5ddc3ae2f9422ca351084e3d5ec41

data/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Brickell Research
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

data/README.md

@@ -0,0 +1,23 @@
+# Token Lens
+
+Basically a combination of [perf](https://perfwiki.github.io/main/) plus [flame-graphs](https://www.brendangregg.com/flamegraphs.html) for local [claude-code](https://code.claude.com/docs/en/overview) usage.
+
+## Architecture
+
+```
+Record Data --> Interpret Data --> Render Data
+```
+
+## Data Sources
+
+1. Tap OTEL console output
+2. Tail session JSONL
+3. Local API proxy
+
+## Contributing
+
+Checkout [CONTRIBUTING.md](./CONTRIBUTING.md).
+
+## License
+
+[MIT](./LICENSE)

data/bin/token-lens

@@ -0,0 +1,7 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require "rubygems"
+require "bundler/setup"
+require "token_lens/cli"
+TokenLens::CLI.start(ARGV)

data/lib/token_lens/cli.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+require "thor"
+require "token_lens/commands/record"
+require "token_lens/commands/render"
+
+module TokenLens
+  class CLI < Thor
+    desc "record", "Tail the active session and capture events to stdout"
+    option :duration_in_seconds, type: :numeric, default: 30, desc: "Seconds to record"
+    def record
+      Commands::Record.new(duration_in_seconds: options[:duration_in_seconds]).run
+    end
+
+    desc "render", "Render a captured session as a flame graph"
+    option :file_path, type: :string, required: true, desc: "Path to the captured JSON file"
+    option :output, type: :string, default: "flame.html", desc: "Output HTML path"
+    def render
+      Commands::Render.new(file_path: options[:file_path], output: options[:output]).run
+    end
+  end
+end

data/lib/token_lens/commands/record.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+require "json"
+require "token_lens/sources/jsonl"
+
+module TokenLens
+  module Commands
+    class Record
+      def initialize(duration_in_seconds:)
+        @duration_in_seconds = duration_in_seconds
+      end
+
+      def run
+        warn "Recording for #{@duration_in_seconds}s... (Ctrl+C to stop early)"
+
+        queue = Queue.new
+        events = []
+
+        thread = Thread.new { Sources::Jsonl.new(queue).start }
+        drain_thread = Thread.new { loop { events << queue.pop } }
+
+        trap("INT") { finish(thread, drain_thread, queue, events) }
+        trap("TERM") { finish(thread, drain_thread, queue, events) }
+
+        sleep @duration_in_seconds
+        finish(thread, drain_thread, queue, events)
+      end
+
+      private
+
+      def finish(thread, drain_thread, queue, events)
+        thread.kill
+        drain_thread.kill
+        events << queue.pop until queue.empty?
+        warn "\nCaptured #{events.size} events"
+        $stdout.puts JSON.generate(events)
+        exit 0
+      end
+    end
+  end
+end

data/lib/token_lens/commands/render.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+require "token_lens/parser"
+require "token_lens/renderer/reshaper"
+require "token_lens/renderer/annotator"
+require "token_lens/renderer/layout"
+require "token_lens/renderer/html"
+
+module TokenLens
+  module Commands
+    class Render
+      def initialize(file_path:, output:)
+        @file_path = file_path
+        @output = output
+      end
+
+      def run
+        tree = Parser.new(file_path: @file_path).parse
+        tree = Renderer::Reshaper.new.reshape(tree)
+        Renderer::Annotator.new.annotate(tree)
+        Renderer::Layout.new.layout(tree)
+        html = Renderer::Html.new.render(tree)
+        File.write(@output, html)
+        warn "Wrote #{@output}"
+      end
+    end
+  end
+end

data/lib/token_lens/parser.rb

@@ -0,0 +1,171 @@
+# frozen_string_literal: true
+
+require "json"
+require "token_lens/tokens/jsonl"
+
+module TokenLens
+  class ParseError < StandardError; end
+
+  class Parser
+    def initialize(file_path:)
+      @file_path = file_path
+    end
+
+    def parse
+      raw_events = JSON.parse(read_file).map { |e| e["event"] }
+      tokens = raw_events
+        .map { |e| Tokens::Jsonl.from_raw(e) }
+        .select { |t| t.type == "user" || t.type == "assistant" }
+      tree = build_tree(tokens, raw_events)
+      attach_subagent_turns(tree, raw_events)
+      attach_task_notifications(tree)
+      tree
+    end
+
+    private
+
+    def build_tree(tokens, raw_events = [])
+      index = tokens.each_with_object({}) { |t, h| h[t.uuid] = {token: t, children: []} }
+
+      # Build a parent map covering ALL raw events (including filtered progress/tool-result
+      # events) so we can walk through gaps when a token's direct parent was filtered out.
+      raw_parent = {}
+      raw_events.each { |e| raw_parent[e["uuid"]] = e["parentUuid"] if e["uuid"] }
+
+      roots = []
+      index.each_value do |node|
+        parent_uuid = node[:token].parent_uuid
+        # Walk up through filtered events to find the nearest indexed ancestor.
+        hops = 0
+        while parent_uuid && !index[parent_uuid] && hops < 20
+          parent_uuid = raw_parent[parent_uuid]
+          hops += 1
+        end
+        parent = parent_uuid && index[parent_uuid]
+        parent ? parent[:children] << node : roots << node
+      end
+
+      roots
+    end
+
+    # Extract subagent turns from agent_progress events and attach them as
+    # is_sidechain children of the assistant turn that invoked the Agent tool.
+    def attach_subagent_turns(tree, raw_events)
+      progress_by_tool_use = Hash.new { |h, k| h[k] = [] }
+      raw_events.each do |evt|
+        next unless evt["type"] == "progress"
+        next unless evt.dig("data", "type") == "agent_progress"
+        next unless evt.dig("data", "message", "type") == "assistant"
+        tool_use_id = evt["parentToolUseID"]
+        next unless tool_use_id
+        progress_by_tool_use[tool_use_id] << evt
+      end
+
+      return if progress_by_tool_use.empty?
+
+      # Index all nodes by their tool_use content IDs
+      tool_use_node = {}
+      flatten_nodes(tree).each do |node|
+        node[:token].tool_uses.each { |tu| tool_use_node[tu["id"]] = node }
+      end
+
+      progress_by_tool_use.each do |tool_use_id, evts|
+        parent = tool_use_node[tool_use_id]
+        next unless parent
+        parent[:children] += build_subagent_nodes(evts)
+      end
+    end
+
+    # Collapse streaming chains (same requestId = one API call) and build tokens.
+    def build_subagent_nodes(events)
+      by_request = Hash.new { |h, k| h[k] = [] }
+      events.each do |evt|
+        req_id = evt.dig("data", "message", "requestId") || evt["uuid"]
+        by_request[req_id] << evt
+      end
+
+      by_request.values
+        .sort_by { |g| g.first["timestamp"] || "" }
+        .map { |group| subagent_token(group) }
+        .map { |t| {token: t, children: []} }
+    end
+
+    def subagent_token(group)
+      representative = group.first
+      msg_data = representative.dig("data", "message")
+      inner = msg_data["message"] || {}
+      usage = inner["usage"] || {}
+
+      # Combine tool_uses across streaming events in this API call (parallel tools)
+      combined_tool_uses = group.flat_map { |evt|
+        Array(evt.dig("data", "message", "message", "content"))
+          .select { |b| b.is_a?(Hash) && b["type"] == "tool_use" }
+      }.uniq { |tu| tu["id"] }
+
+      content = combined_tool_uses.empty? ? Array(inner["content"]) : combined_tool_uses
+
+      Tokens::Jsonl.new(
+        uuid: msg_data["uuid"] || representative["uuid"],
+        parent_uuid: nil,
+        request_id: msg_data["requestId"],
+        type: "assistant",
+        role: "assistant",
+        model: inner["model"],
+        is_sidechain: true,
+        agent_id: representative.dig("data", "agentId"),
+        content: content,
+        input_tokens: usage["input_tokens"].to_i,
+        output_tokens: usage["output_tokens"].to_i,
+        cache_read_tokens: usage["cache_read_input_tokens"].to_i,
+        cache_creation_tokens: usage["cache_creation_input_tokens"].to_i,
+        marginal_input_tokens: 0,
+        timestamp: representative["timestamp"],
+        is_compaction: false
+      )
+    end
+
+    # Wire task-notification user turns back to the Agent call that spawned them.
+    # Each <task-notification> contains a <tool-use-id> that matches an Agent
+    # tool call in the main thread. We detach the notification from its current
+    # tree position, mark it is_sidechain, and attach it under the Agent call node.
+    def attach_task_notifications(tree)
+      all = flatten_nodes(tree)
+
+      agent_node_by_tool_use = {}
+      all.each do |node|
+        node[:token].tool_uses.each do |tu|
+          next unless tu["name"] == "Agent"
+          agent_node_by_tool_use[tu["id"]] = node
+        end
+      end
+      return if agent_node_by_tool_use.empty?
+
+      all.select { |n| n[:token].is_task_notification? }.each do |node|
+        tool_use_id = node[:token].human_text
+          .match(/<tool-use-id>\s*(.*?)\s*<\/tool-use-id>/m)&.[](1)
+        next unless tool_use_id
+        agent_node = agent_node_by_tool_use[tool_use_id]
+        next unless agent_node
+        next unless remove_node(tree, node)
+
+        node[:token] = node[:token].with(is_sidechain: true)
+        agent_node[:children] << node
+      end
+    end
+
+    def remove_node(nodes, target)
+      return true if nodes.delete(target)
+      nodes.any? { |node| remove_node(node[:children], target) }
+    end
+
+    def flatten_nodes(nodes)
+      nodes.flat_map { |n| [n, *flatten_nodes(n[:children])] }
+    end
+
+    def read_file
+      File.read(@file_path)
+    rescue => e
+      raise TokenLens::ParseError, "Failed to read file: #{e.message}"
+    end
+  end
+end

data/lib/token_lens/pricing.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module TokenLens
+  module Pricing
+    # Prices in USD per million tokens. Source: platform.claude.com/docs/en/about-claude/pricing
+    # Last verified: 2026-03-23
+    #
+    # cache_creation = 5-minute cache write (1.25x input). The API reports this as
+    # cache_creation_input_tokens in the usage object.
+    # cache_read    = cache hit (0.1x input).
+    #
+    # Entries are matched via String#start_with? in order — put more specific prefixes first.
+    TABLE = {
+      # --- Opus 4.5 / 4.6 (new pricing tier: $5/$25) ---
+      "claude-opus-4-6" => {input: 5.0, cache_read: 0.50, cache_creation: 6.25, output: 25.0},
+      "claude-opus-4-5" => {input: 5.0, cache_read: 0.50, cache_creation: 6.25, output: 25.0},
+
+      # --- Opus 4.0 / 4.1 (original tier: $15/$75) ---
+      "claude-opus-4" => {input: 15.0, cache_read: 1.50, cache_creation: 18.75, output: 75.0},
+
+      # --- Sonnet 4.x (all variants same price) ---
+      "claude-sonnet-4" => {input: 3.0, cache_read: 0.30, cache_creation: 3.75, output: 15.0},
+
+      # --- Haiku 4.5 ---
+      "claude-haiku-4-5" => {input: 1.0, cache_read: 0.10, cache_creation: 1.25, output: 5.0},
+
+      # --- Haiku 4.x fallback ---
+      "claude-haiku-4" => {input: 1.0, cache_read: 0.10, cache_creation: 1.25, output: 5.0},
+
+      # --- Claude 3.x (legacy, new-style IDs like claude-sonnet-3-7) ---
+      "claude-sonnet-3" => {input: 3.0, cache_read: 0.30, cache_creation: 3.75, output: 15.0},
+      "claude-haiku-3-5" => {input: 0.80, cache_read: 0.08, cache_creation: 1.00, output: 4.0},
+
+      # --- Claude 3.x (old-style IDs like claude-3-opus-20240229) ---
+      "claude-3-opus" => {input: 15.0, cache_read: 1.50, cache_creation: 18.75, output: 75.0},
+      "claude-3-5-sonnet" => {input: 3.0, cache_read: 0.30, cache_creation: 3.75, output: 15.0},
+      "claude-3-sonnet" => {input: 3.0, cache_read: 0.30, cache_creation: 3.75, output: 15.0},
+      "claude-3-5-haiku" => {input: 0.80, cache_read: 0.08, cache_creation: 1.00, output: 4.0},
+      "claude-3-haiku" => {input: 0.25, cache_read: 0.03, cache_creation: 0.30, output: 1.25}
+    }.freeze
+
+    # Fallback when model string is nil or unrecognised — use Sonnet 4 rates
+    FALLBACK = {input: 3.0, cache_read: 0.30, cache_creation: 3.75, output: 15.0}.freeze
+
+    def self.for_model(model)
+      return FALLBACK unless model
+      _prefix, rates = TABLE.find { |prefix, _| model.start_with?(prefix) }
+      rates || FALLBACK
+    end
+  end
+end

data/lib/token_lens/renderer/annotator.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module TokenLens
+  module Renderer
+    class Annotator
+      def annotate(nodes, depth = 0)
+        nodes.each do |node|
+          node[:depth] = depth
+          annotate(node[:children], depth + 1)
+          child_tokens = node[:children].sum { |c| c[:subtree_tokens] }
+          child_cost = node[:children].sum { |c| c[:subtree_cost] }
+          node[:subtree_tokens] = [node[:token].display_width, 1].max + child_tokens
+          node[:subtree_cost] = node[:token].cost_usd + child_cost
+        end
+        nodes
+      end
+    end
+  end
+end