token-lens 0.1.0

data/lib/token_lens/renderer/layout.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+module TokenLens
+  module Renderer
+    class Layout
+      CANVAS_WIDTH = 1200
+      ROW_HEIGHT = 32
+
+      def initialize(canvas_width: CANVAS_WIDTH)
+        @canvas_width = canvas_width
+      end
+
+      def layout(nodes)
+        max_depth = all_nodes(nodes).map { |n| n[:depth] }.max || 0
+        total = nodes.sum { |n| n[:subtree_tokens] }
+        scale = (total > 0) ? @canvas_width.to_f / total : 1.0
+        position(nodes, x: 0, scale: scale, max_depth: max_depth)
+
+        total_cost = nodes.sum { |n| n[:subtree_cost] }
+        cost_scale = (total_cost > 0) ? @canvas_width.to_f / total_cost : 1.0
+        position_cost(nodes, x: 0, scale: cost_scale, max_depth: max_depth)
+
+        nodes
+      end
+
+      private
+
+      # Bottom-up layout: roots at bottom (y = max_depth * ROW_HEIGHT),
+      # deepest children at top (y = 0).
+      def position(nodes, x:, scale:, max_depth:)
+        cursor = x
+        nodes.each do |node|
+          node[:x] = cursor
+          node[:y] = (max_depth - node[:depth]) * ROW_HEIGHT
+          node[:w] = (node[:subtree_tokens] * scale).round
+          position(node[:children], x: cursor, scale: scale, max_depth: max_depth)
+          cursor += node[:w]
+        end
+      end
+
+      def position_cost(nodes, x:, scale:, max_depth:)
+        cursor = x
+        nodes.each do |node|
+          node[:cost_x] = cursor
+          node[:cost_w] = (node[:subtree_cost] * scale).round
+          position_cost(node[:children], x: cursor, scale: scale, max_depth: max_depth)
+          cursor += node[:cost_w]
+        end
+      end
+
+      def all_nodes(nodes)
+        nodes.flat_map { |n| [n, *all_nodes(n[:children])] }
+      end
+    end
+  end
+end

data/lib/token_lens/renderer/reshaper.rb

@@ -0,0 +1,88 @@
+# frozen_string_literal: true
+
+module TokenLens
+  module Renderer
+    class Reshaper
+      def reshape(nodes)
+        nodes = collapse_streaming(nodes)
+        nodes.flat_map { |node| process_root(node) }
+      end
+
+      private
+
+      # Collapse streaming chains: thinking → text → tool_use events emitted by
+      # Claude Code for a single API response, detected by identical input usage.
+      def collapse_streaming(nodes)
+        nodes.flat_map do |node|
+          if streaming_intermediate?(node)
+            collapse_streaming(node[:children])
+          else
+            node[:children] = collapse_streaming(node[:children])
+            [node]
+          end
+        end
+      end
+
+      def streaming_intermediate?(node)
+        return false unless node[:token].role == "assistant"
+        return false unless node[:children].size == 1
+        child = node[:children].first
+        return false unless child[:token].role == "assistant"
+        t, c = node[:token], child[:token]
+        # Prefer request_id equality (same API call); fall back to token count fingerprint
+        if t.request_id && c.request_id
+          t.request_id == c.request_id
+        else
+          t.input_tokens == c.input_tokens &&
+            t.cache_read_tokens == c.cache_read_tokens &&
+            t.cache_creation_tokens == c.cache_creation_tokens
+        end
+      end
+
+      # Re-root the tree around human prompt nodes. Human prompts become roots;
+      # the linear assistant chain beneath them becomes a flat list of siblings.
+      def process_root(node)
+        t = node[:token]
+        if t.is_human_prompt?
+          siblings = flatten_thread(node[:children], prev_input: 0)
+          [node.merge(children: siblings)]
+        elsif t.role == "user"
+          # Tool-result-only user at root level — hoist children
+          node[:children].flat_map { |c| process_root(c) }
+        else
+          # Orphan assistant root (no human prompt ancestor)
+          flatten_thread([node], prev_input: 0)
+        end
+      end
+
+      # Flatten a linear user→assistant→user(tool_result)→assistant chain into
+      # a flat list of assistant siblings, computing marginal_input_tokens deltas.
+      # Sidechain children stay nested under the assistant that spawned them.
+      def flatten_thread(nodes, prev_input:)
+        nodes.flat_map do |node|
+          t = node[:token]
+          if t.role == "user" && !t.is_human_prompt?
+            flatten_thread(node[:children], prev_input: prev_input)
+          elsif t.role == "assistant"
+            marginal = [t.input_tokens - prev_input, 0].max
+            compaction = prev_input > 0 && t.input_tokens < prev_input * 0.5
+            sidechain = node[:children].select { |c| c[:token].is_sidechain }
+            chain = node[:children].reject { |c| c[:token].is_sidechain }
+            # Flatten the response chain inside task-notification sidechains so
+            # they don't create arbitrarily deep linked-list nesting.
+            sidechain = sidechain.map do |sc|
+              sc[:token].is_task_notification? ? sc.merge(children: flatten_thread(sc[:children], prev_input: 0)) : sc
+            end
+            updated = node.merge(
+              token: t.with(marginal_input_tokens: marginal, is_compaction: compaction),
+              children: sidechain
+            )
+            [updated] + flatten_thread(chain, prev_input: t.input_tokens)
+          else
+            process_root(node)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/token_lens/session.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+require "json"
+require "pathname"
+
+module TokenLens
+  module Session
+    # Claude Code stores sessions at ~/.claude/projects/<encoded-cwd>/<session-id>.jsonl
+    # where encoded-cwd is the absolute working directory with every non-alphanumeric
+    # character replaced by a hyphen (e.g. /Users/me/proj => -Users-me-proj)
+    CLAUDE_DIR = Pathname.new(File.expand_path("~/.claude/projects"))
+
+    def self.encoded_cwd(dir = Dir.pwd)
+      dir.gsub(/[^a-zA-Z0-9]/, "-")
+    end
+
+    def self.active_jsonl(dir = Dir.pwd)
+      project_dir = CLAUDE_DIR / encoded_cwd(dir)
+
+      # find all *.jsonl files in the project directory
+      jsonl_files = project_dir.glob("*.jsonl")
+      raise "No session files found in #{project_dir}" if jsonl_files.empty?
+
+      # get most recently modified file
+      jsonl_files.max_by(&:mtime)
+    end
+
+    def self.tail(path, &block)
+      last_pos = File.size(path)
+      loop do
+        sleep 0.1
+        current_size = File.size(path)
+        # if file size hasn't changed, skip
+        next if current_size == last_pos
+
+        # otherwise, read new lines from the file
+        File.open(path) do |f|
+          f.seek(last_pos)
+          f.each_line { |line| block.call(JSON.parse(line)) }
+          last_pos = f.pos
+        end
+      end
+    end
+  end
+end

data/lib/token_lens/sources/jsonl.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+require "json"
+require "token_lens/session"
+
+module TokenLens
+  module Sources
+    class Jsonl
+      def initialize(queue)
+        @queue = queue
+        @path = Session.active_jsonl
+      end
+
+      def start
+        warn "  [jsonl] tailing #{@path.basename}"
+        last_pos = File.size(@path)
+
+        loop do
+          sleep 0.1
+          current_size = File.size(@path)
+          next if current_size == last_pos
+
+          File.open(@path) do |f|
+            f.seek(last_pos)
+            f.each_line do |line|
+              event = JSON.parse(line)
+              @queue << {source: "jsonl", event: event}
+            end
+            last_pos = f.pos
+          end
+        end
+      end
+    end
+  end
+end

data/lib/token_lens/tokens/jsonl.rb

@@ -0,0 +1,99 @@
+# frozen_string_literal: true
+
+require_relative "../pricing"
+
+module TokenLens
+  module Tokens
+    Jsonl = Data.define(
+      :uuid,
+      :parent_uuid,
+      :request_id,
+      :type,
+      :role,
+      :model,
+      :is_sidechain,
+      :agent_id,
+      :content,
+      :input_tokens,
+      :output_tokens,
+      :cache_read_tokens,
+      :cache_creation_tokens,
+      :marginal_input_tokens,
+      :timestamp,
+      :is_compaction
+    ) do
+      def self.from_raw(raw)
+        msg = raw["message"] || {}
+        usage = msg["usage"] || {}
+
+        new(
+          uuid: raw["uuid"],
+          parent_uuid: raw["parentUuid"],
+          request_id: raw["requestId"],
+          type: raw["type"],
+          role: msg["role"],
+          model: msg["model"],
+          is_sidechain: raw["isSidechain"] || false,
+          agent_id: nil,
+          content: Array(msg["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: raw["timestamp"],
+          is_compaction: false
+        )
+      end
+
+      def cost_usd
+        p = Pricing.for_model(model)
+        (marginal_input_tokens * p[:input] +
+          cache_read_tokens * p[:cache_read] +
+          cache_creation_tokens * p[:cache_creation] +
+          output_tokens * p[:output]) / 1_000_000.0
+      end
+
+      def total_tokens
+        input_tokens + output_tokens + cache_read_tokens + cache_creation_tokens
+      end
+
+      def display_width
+        marginal_input_tokens + cache_creation_tokens + output_tokens
+      end
+
+      def assistant?
+        role == "assistant"
+      end
+
+      def is_human_prompt?
+        return false unless role == "user"
+        return false if tool_results.any?
+        content.any? { |b| b.is_a?(String) || (b.is_a?(Hash) && b["type"] == "text") }
+      end
+
+      def is_task_notification?
+        is_human_prompt? && human_text.start_with?("<task-notification>")
+      end
+
+      def task_notification_summary
+        human_text.match(/<summary>(.*?)<\/summary>/m)&.[](1)&.strip
+      end
+
+      def human_text
+        block = content.find { |b| b.is_a?(String) }
+        return block if block
+        block = content.find { |b| b.is_a?(Hash) && b["type"] == "text" }
+        block&.dig("text") || ""
+      end
+
+      def tool_uses
+        content.select { |b| b.is_a?(Hash) && b["type"] == "tool_use" }
+      end
+
+      def tool_results
+        content.select { |b| b.is_a?(Hash) && b["type"] == "tool_result" }
+      end
+    end
+  end
+end

data/lib/token_lens/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module TokenLens
+  VERSION = "0.1.0"
+end

data/lib/token_lens.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+require_relative "token_lens/version"
+
+module TokenLens
+  def self.hello
+    "token-lens v#{VERSION}"
+  end
+end

metadata

@@ -0,0 +1,71 @@
+--- !ruby/object:Gem::Specification
+name: token-lens
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- rob durst
+bindir: bin
+cert_chain: []
+date: 2026-03-23 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: thor
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.3'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.3'
+email:
+- me@robdurst.com
+executables:
+- token-lens
+extensions: []
+extra_rdoc_files: []
+files:
+- LICENSE
+- README.md
+- bin/token-lens
+- lib/token_lens.rb
+- lib/token_lens/cli.rb
+- lib/token_lens/commands/record.rb
+- lib/token_lens/commands/render.rb
+- lib/token_lens/parser.rb
+- lib/token_lens/pricing.rb
+- lib/token_lens/renderer/annotator.rb
+- lib/token_lens/renderer/html.rb
+- lib/token_lens/renderer/layout.rb
+- lib/token_lens/renderer/reshaper.rb
+- lib/token_lens/session.rb
+- lib/token_lens/sources/jsonl.rb
+- lib/token_lens/tokens/jsonl.rb
+- lib/token_lens/version.rb
+homepage: https://github.com/BrickellResearch/token-lens
+licenses:
+- MIT
+metadata: {}
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '3.2'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.6.2
+specification_version: 4
+summary: Flame graphs for Claude Code token usage
+test_files: []