token-lens 0.3.0 → 0.5.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 5df4aa9d712d04ae90513777b2b34b541bb3b77038fadd637725f1bf1599db3f
-  data.tar.gz: 52f773f6c7bfd677a29bef81ba14bb23e2b01ec54071468ee6ba7be2e4632b9a
+  metadata.gz: d15dc78e85a94ff7c3e4c22388e74a9084543cbe00936dcc26c155f2c1e07191
+  data.tar.gz: e66a0b7ba2753b5129f76c5900a8c81dcb6b51bdfb2a17b34e4a9d4c9925a7cc
 SHA512:
-  metadata.gz: 0b271f96c07b9c74db5ca42e4ffff7a7259e5a63893e949a9ab0fc3dd2bc4761df0c783151769a42cd320dc6e5574ae408a5b2dbfe2133cc839bab60d12f638a
-  data.tar.gz: 8f296614fd7f957189e07961223fea18a2242fee5f090860a2b49d6f69daf90f92dadc87a289ce4b22cbaf739d0d25cfdfbd2c2cac062fe4ed7e23f952db156c
+  metadata.gz: 9f5cdbfe0e769a8cac9c5d610336e66d1a8f9d07b0767bd31d3250fbe47425597eb89f5bc9e04e7a3ab8f97e2effe7c90d073f3fd8e02236ca4b210d6e81f6a3
+  data.tar.gz: 3fe4f1b44190dd712cf6c4353c3155fb0684a3192e987436a26611005a9775be23a09d0ddd34c2dc2845e678435ce0f2842068b335b245dadef994d17a3c009b

data/README.md

@@ -1,5 +1,12 @@
 # Token Lens
 
+[![Gem Version](https://img.shields.io/gem/v/token-lens)](https://rubygems.org/gems/token-lens)
+[![Gem Downloads](https://img.shields.io/gem/dt/token-lens)](https://rubygems.org/gems/token-lens)
+[![CI](https://github.com/Brickell-Research/token-lens/actions/workflows/ci.yml/badge.svg)](https://github.com/Brickell-Research/token-lens/actions/workflows/ci.yml)
+[![Ruby >= 3.2](https://img.shields.io/badge/Ruby-%3E%3D3.2-red?logo=ruby&logoColor=white)](https://www.ruby-lang.org/)
+[![Standard](https://img.shields.io/badge/code_style-standard-brightgreen.svg)](https://github.com/standardrb/standard)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+
 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
@@ -16,19 +23,44 @@ Record Data --> Interpret Data --> Render Data
 
 ## Quick Start
 
+### Zero setup — render any session right now
+
+No recording needed. token-lens reads directly from Claude Code's session files:
 
-**Locally**:
 ```
-bin/token-lens record --duration-in-seconds=30 > capture.json
-bin/token-lens render --file-path=capture.json
+gem install token-lens
+token-lens render
 open flame.html
 ```
 
-**Via gem**:
+That's it. `render` with no arguments finds your most recent Claude Code session and renders it as a flame graph — no prior setup, no capture file, no extra terminal.
+
+### Record a live session
+
+If you want to capture a bounded window while you work (useful for comparing before/after):
+
 ```
-gem install token-lens
-token-lens record --duration-in-seconds=30 > capture.json
-token-lens render --file-path=capture.json
+token-lens record --duration-in-seconds=60
+# ... do your Claude Code work in another terminal ...
+token-lens render
+open flame.html
+```
+
+Captures auto-save to `~/.token-lens/sessions/<timestamp>.json`. `render` always picks the most recent capture first, then falls back to the live session JSONL if no captures exist.
+
+### Options
+
+```
+token-lens record --duration-in-seconds=300   # record for 5 minutes
+token-lens record --output=my-session.json    # save to a specific path
+token-lens render --file-path=my-session.json # render a specific capture
+token-lens render --output=report.html        # write HTML to a custom path
+```
+
+### Locally
+
+```
+bin/token-lens render
 open flame.html
 ```
 

data/lib/token_lens/cli.rb

@@ -6,18 +6,20 @@ require "token_lens/commands/render"
 
 module TokenLens
   class CLI < Thor
-    desc "record", "Tail the active session and capture events to stdout"
+    desc "record", "Tail the active session and auto-save a capture file"
     option :duration_in_seconds, type: :numeric, default: 30, desc: "Seconds to record"
     option :project_dir, type: :string, desc: "Working directory of the Claude Code session to record (default: auto-detect)"
+    option :output, type: :string, desc: "Save path for the capture (default: ~/.token-lens/sessions/<timestamp>.json)"
     def record
       Commands::Record.new(
         duration_in_seconds: options[:duration_in_seconds],
-        project_dir: options[:project_dir]
+        project_dir: options[:project_dir],
+        output: options[:output]
       ).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 :file_path, type: :string, desc: "Path to the captured JSON file (default: most recent session)"
     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

data/lib/token_lens/commands/record.rb

@@ -1,14 +1,18 @@
 # frozen_string_literal: true
 
 require "json"
+require "fileutils"
 require "token_lens/sources/jsonl"
 
 module TokenLens
   module Commands
     class Record
-      def initialize(duration_in_seconds:, project_dir: nil)
+      SESSIONS_DIR = Pathname.new(Dir.home).join(".token-lens", "sessions")
+
+      def initialize(duration_in_seconds:, project_dir: nil, output: nil)
         @duration_in_seconds = duration_in_seconds
         @project_dir = project_dir
+        @output = output
       end
 
       def run
@@ -34,9 +38,18 @@ module TokenLens
         drain_thread.kill
         events << queue.pop until queue.empty?
         warn "\nCaptured #{events.size} events"
-        $stdout.puts JSON.generate(events)
+        path = save_path
+        FileUtils.mkdir_p(path.dirname)
+        path.write(JSON.generate(events))
+        warn "Saved to #{path}"
         exit 0
       end
+
+      def save_path
+        return Pathname.new(@output) if @output
+        timestamp = Time.now.strftime("%Y-%m-%d_%H-%M-%S")
+        SESSIONS_DIR.join("#{timestamp}.json")
+      end
     end
   end
 end

data/lib/token_lens/commands/render.rb

@@ -5,24 +5,39 @@ require "token_lens/renderer/reshaper"
 require "token_lens/renderer/annotator"
 require "token_lens/renderer/layout"
 require "token_lens/renderer/html"
+require "token_lens/session"
 
 module TokenLens
   module Commands
     class Render
-      def initialize(file_path:, output:)
+      def initialize(output:, file_path: nil)
         @file_path = file_path
         @output = output
       end
 
       def run
-        tree = Parser.new(file_path: @file_path).parse
+        path = resolve_path
+        warn "Rendering #{path}"
+        tree = Parser.new(file_path: path).parse
         tree = Renderer::Reshaper.new.reshape(tree)
+        tree.sort_by! { |n| n[:token].timestamp || "" }
         Renderer::Annotator.new.annotate(tree)
-        Renderer::Layout.new.layout(tree)
-        html = Renderer::Html.new.render(tree)
+        canvas_width = Renderer::Layout.new.layout(tree)
+        html = Renderer::Html.new(canvas_width: canvas_width).render(tree)
         File.write(@output, html)
         warn "Wrote #{@output}"
       end
+
+      private
+
+      def resolve_path
+        return @file_path if @file_path
+        sessions = Pathname.new(Dir.home).join(".token-lens", "sessions")
+        saved = sessions.glob("*.json").max_by(&:mtime)
+        return saved if saved
+        warn "No saved captures found — reading active Claude Code session directly"
+        Session.active_or_latest_jsonl
+      end
     end
   end
 end

data/lib/token_lens/parser.rb

@@ -12,7 +12,7 @@ module TokenLens
     end
 
     def parse
-      raw_events = JSON.parse(read_file).map { |e| e["event"] }
+      raw_events = read_raw_events
       tokens = raw_events
         .map { |e| Tokens::Jsonl.from_raw(e) }
         .select { |t| t.type == "user" || t.type == "assistant" }
@@ -162,6 +162,25 @@ module TokenLens
       nodes.flat_map { |n| [n, *flatten_nodes(n[:children])] }
     end
 
+    def read_raw_events
+      content = read_file
+      if content.lstrip.start_with?("[")
+        # Captured format: JSON array of {"event": {...}} wrappers produced by `record`
+        JSON.parse(content).map { |e| e["event"] }
+      else
+        # Raw JSONL: newline-delimited events from ~/.claude/projects/*/...jsonl
+        content.each_line.filter_map { |line|
+          begin
+            JSON.parse(line)
+          rescue
+            nil
+          end
+        }
+      end
+    rescue => e
+      raise TokenLens::ParseError, "Failed to parse #{@file_path}: #{e.message}"
+    end
+
     def read_file
       File.read(@file_path)
     rescue => e

data/lib/token_lens/renderer/html.rb

@@ -6,16 +6,21 @@ module TokenLens
   module Renderer
     class Html
       ROW_HEIGHT = 32
-      COORD_WIDTH = 1200  # logical coordinate space (matches Layout)
+      COORD_WIDTH = 1200  # default logical coordinate space
       FONT_SIZE = 13
-      MIN_LABEL_PCT = 3.0  # hide label if bar is narrower than 3% of canvas
+      MIN_LABEL_PX = 60  # hide label if bar is narrower than 60px
+
+      def initialize(canvas_width: COORD_WIDTH)
+        @canvas_width = canvas_width
+      end
 
       LEGEND_ITEMS = [
         ["bar-c-human", "User prompt"],
         ["bar-c-task", "Task callback"],
         ["bar-c-assistant", "Assistant response"],
         ["bar-c-tool", "Tool call"],
-        ["bar-c-sidechain", "Subagent turn"]
+        ["bar-c-sidechain", "Subagent turn"],
+        ["bar-compaction", "Compaction"]
       ].freeze
 
       CONTEXT_LIMIT = 200_000  # all current Claude models
@@ -28,7 +33,7 @@ module TokenLens
         total_top = (max_depth + 1) * ROW_HEIGHT
         total_tokens = nodes.sum { |n| n[:subtree_tokens] }
         total_cost = nodes.sum { |n| n[:subtree_cost] }
-        total_tip = escape_js(escape_html(total_summary(all)))
+        total_tip = escape_js(escape_html(total_summary(all_flat)))
         @reread_files = build_reread_map(all)
         @thread_count = nodes.length
         @thread_numbers = {}
@@ -39,6 +44,8 @@ module TokenLens
           next unless id && !@agent_labels.key?(id)
           @agent_labels[id] = "A#{@agent_labels.size + 1}"
         end
+        assign_alternation(nodes)
+        @hm_count = nodes.length  # overridden by heatmap_html after grouping
         token_total_lbl = "TOTAL · #{fmt(total_tokens)} tokens"
         cost_total_lbl = "TOTAL · #{fmt_cost(total_cost)}"
 
@@ -56,23 +63,27 @@ module TokenLens
           <body>
           <div class="header">
             <div>
-              <div class="summary">#{summary_text(all)}</div>
-              <div class="legend">#{legend_html}</div>
+              <div class="summary">#{summary_text(all_flat)}</div>
+              <div class="legend" id="legend" style="display:none">#{legend_html}</div>
             </div>
             <div class="header-btns">
               <button class="theme-btn" id="theme-btn" onclick="toggleTheme()">&#x25D0; Light</button>
               <button class="summary-btn" id="summary-btn" onclick="toggleSummary()">&#x2261; Summary</button>
               <button class="cost-btn" id="cost-btn" onclick="toggleCostView()">$ Cost view</button>
-              <button class="reset-btn" id="reset-btn" onclick="unzoom()">&#x21A9; Reset zoom</button>
+              <button class="reset-btn" id="reset-btn" onclick="resetZoom()">&#x21A9; Reset zoom</button>
+              <button class="export-btn" onclick="exportSVG()">&#x2913; Export</button>
             </div>
           </div>
-          <div class="spacer"></div>
-          <div class="flame" style="height:#{flame_height}px">
-          <div class="bar total-bar" style="left:0%;width:100%;top:#{total_top}px" data-ox="0" data-ow="#{COORD_WIDTH}" data-cx="0" data-cw="#{COORD_WIDTH}" onmouseover="tip('#{total_tip}')" onmouseout="tip('')" onclick="unzoom()"><span class="lbl total-lbl" id="total-lbl" data-token-text="#{token_total_lbl}" data-cost-text="#{cost_total_lbl}">#{token_total_lbl}</span></div>
+          #{heatmap_html(nodes)}
+          <div id="hm-back" class="hm-back" style="display:none" onclick="closePrompt()">&#x2190; All prompts</div>
+          <div class="flame-wrap" id="flame-wrap" style="display:none">
+          <div class="flame" style="width:#{@canvas_width}px;height:#{flame_height}px">
+          <div class="bar total-bar" style="left:0%;width:100%;top:#{total_top}px" data-ox="0" data-ow="#{@canvas_width}" data-cx="0" data-cw="#{@canvas_width}" onmouseover="tip('#{total_tip}')" onmouseout="tip('')" onclick="if(hmActiveIdx<0)unzoom()"><span class="lbl total-lbl" id="total-lbl" data-token-text="#{token_total_lbl}" data-cost-text="#{cost_total_lbl}">#{token_total_lbl}</span></div>
           #{all.map { |n| bar_html(n) }.join("\n")}
           </div>
+          </div>
           <div id="ftip" class="floattip"></div>
-          <div class="tip" id="tip">&nbsp;</div>
+          <div id="tip" class="tip"><span class="tip-label">Hover for details &middot; Click to zoom</span></div>
           #{session_summary_html(all_flat)}
           <script>
           #{js}
@@ -89,7 +100,7 @@ module TokenLens
       end
 
       def pct(val)
-        (val.to_f / COORD_WIDTH * 100).round(4)
+        (val.to_f / @canvas_width * 100).round(4)
       end
 
       def bar_html(node)
@@ -99,10 +110,11 @@ module TokenLens
         tip = escape_html(tooltip(node))
         left = pct(node[:x])
         width = pct(node[:w])
-        lbl_hidden = (width < MIN_LABEL_PCT) ? " style=\"display:none\"" : ""
+        lbl_hidden = (node[:w] < MIN_LABEL_PX) ? " style=\"display:none\"" : ""
         extra_class = " #{color_class(node)}"
+        extra_class += " bar-alt" if node[:alt]
         extra_class += " bar-reread" if reread_bar?(node)
-        extra_class += " bar-compaction" if t.is_compaction
+        # is_compaction on assistant turns is for summary counting only; color is set via color_class on the user prompt
         extra_class += " bar-pressure" if !t.is_compaction && context_pressure?(node)
         ftip = escape_html(token_summary(node))
         mouseover = if t.is_task_notification?
@@ -115,7 +127,7 @@ module TokenLens
         end
         badge = reread_bar?(node) ? "<span class=\"warn-badge\">\u26a0</span>" : ""
         <<~HTML.chomp
-          <div class="bar#{extra_class}" style="left:#{left}%;width:#{width}%;top:#{node[:y]}px" data-ox="#{node[:x]}" data-ow="#{node[:w]}" data-cx="#{node[:cost_x]}" data-cw="#{node[:cost_w]}" data-token-lbl="#{lbl}" data-cost-lbl="#{clbl}" data-ftip="#{ftip}" onmouseover="#{mouseover}" onmouseout="tip('')" onclick="zoom(this)"><span class="lbl"#{lbl_hidden}>#{lbl}</span>#{badge}</div>
+          <div class="bar#{extra_class}" style="left:#{left}%;width:calc(#{width}% + 1px);top:#{node[:y]}px" data-ox="#{node[:x]}" data-ow="#{node[:w]}" data-cx="#{node[:cost_x]}" data-cw="#{node[:cost_w]}" data-token-lbl="#{lbl}" data-cost-lbl="#{clbl}" data-ftip="#{ftip}" onmouseover="#{mouseover}" onmouseout="tip('')" onclick="zoom(this)"><span class="lbl"#{lbl_hidden}>#{lbl}</span>#{badge}</div>
         HTML
       end
 
@@ -184,11 +196,19 @@ module TokenLens
         total_input = raw_input + cached + cache_new
         hit_rate = (total_input > 0 && cached > 0) ? (cached.to_f / total_input * 100).round(0).to_i : nil
         parts = []
+        timestamps = all.map { |n| n[:token].timestamp }.compact.sort
+        duration = begin
+          secs = (Time.parse(timestamps.last) - Time.parse(timestamps.first)).to_i
+          fmt_duration(secs) if timestamps.length >= 2 && secs > 0
+        rescue
+          nil
+        end
         parts << "#{@thread_count} threads" if @thread_count&.> 1
         parts << "#{prompts} #{"prompt".then { |w| (prompts == 1) ? w : "#{w}s" }}"
         parts << "#{tasks} #{"task callback".then { |w| (tasks == 1) ? w : "#{w}s" }}" if tasks > 0
         parts << "#{turns} main #{"turn".then { |w| (turns == 1) ? w : "#{w}s" }}"
         parts << "#{sub} subagent #{"turn".then { |w| (sub == 1) ? w : "#{w}s" }}" if sub > 0
+        parts << duration if duration
         parts << "fresh input: #{fmt(marginal)}" if marginal > 0
         parts << "cached input: #{fmt(cached)}" if cached > 0
         parts << "written to cache: #{fmt(cache_new)}" if cache_new > 0
@@ -214,6 +234,24 @@ module TokenLens
         parts.join(" | ")
       end
 
+      def thread_separators(all)
+        roots = all.select { |n| n[:depth] == 0 }
+        return "" if roots.size <= 1
+        # Emit a vertical line at the right edge of each thread (except the last)
+        roots[0..-2].map { |n|
+          right_x = n[:x] + n[:w]
+          pct_val = pct(right_x)
+          %(<div class="thread-sep" style="left:#{pct_val}%"></div>)
+        }.join("\n")
+      end
+
+      def assign_alternation(siblings)
+        siblings.each_with_index do |node, i|
+          node[:alt] = i.odd?
+          assign_alternation(node[:children])
+        end
+      end
+
       def legend_html
         LEGEND_ITEMS.map { |css_class, lbl|
           %(<span class="legend-item"><span class="legend-swatch #{css_class}"></span>#{lbl}</span>)
@@ -224,6 +262,7 @@ module TokenLens
         t = node[:token]
         return "bar-c-task" if t.is_task_notification?
         return "bar-c-sidechain" if t.is_sidechain
+        return "bar-compaction" if t.is_human_prompt? && t.human_text.start_with?("This session is being continued")
         return "bar-c-human" if t.is_human_prompt?
         case t.role
         when "user" then "bar-c-user"
@@ -350,7 +389,7 @@ module TokenLens
         total_input = raw_input + cached + cache_new
         hit_rate = (total_input > 0 && cached > 0) ? (cached.to_f / total_input * 100).round(1) : nil
         total_cost = assistant_nodes.sum { |n| n[:token].cost_usd }
-        compactions = assistant_nodes.count { |n| n[:token].is_compaction }
+        compactions = all.count { |n| n[:token].is_human_prompt? && n[:token].human_text.start_with?("This session is being continued") }
         pressure = assistant_nodes.count { |n| context_pressure?(n) }
         models = assistant_nodes.map { |n| n[:token].model }.compact
           .map { |m| model_short(m) }.uniq.join(", ")
@@ -432,6 +471,96 @@ module TokenLens
         str.gsub("\\") { "\\\\" }.gsub("'") { "\\'" }.gsub("\n", "\\n").gsub("\r", "\\r")
       end
 
+      def heatmap_color(value, min_val, max_val)
+        t = (max_val == min_val) ? 0.5 : ((value - min_val).to_f / (max_val - min_val))
+        t **= 0.7
+        r = (32 + (255 - 32) * t).round.clamp(0, 255)
+        g = (5 + (20 - 5) * t).round.clamp(0, 255)
+        b = (16 + (147 - 16) * t).round.clamp(0, 255)
+        "#%02x%02x%02x" % [r, g, b]
+      end
+
+      def heatmap_color_token(value, min_val, max_val)
+        t = (max_val == min_val) ? 0.5 : ((value - min_val).to_f / (max_val - min_val))
+        t **= 0.7
+        r = (7 + (0 - 7) * t).round.clamp(0, 255)
+        g = (48 + (206 - 48) * t).round.clamp(0, 255)
+        b = (48 + (209 - 48) * t).round.clamp(0, 255)
+        "#%02x%02x%02x" % [r, g, b]
+      end
+
+      def compaction_node?(node)
+        node[:token].is_human_prompt? && node[:token].human_text.start_with?("This session is being continued")
+      end
+
+      def heatmap_html(nodes)
+        # Merge each compaction node into the preceding group — it's overhead from that prompt
+        groups = []
+        nodes.each do |node|
+          if compaction_node?(node) && groups.any?
+            groups.last << node
+          else
+            groups << [node]
+          end
+        end
+        @hm_count = groups.length
+
+        group_tokens = groups.map { |g| g.sum { |n| n[:subtree_tokens] } }
+        group_costs = groups.map { |g| g.sum { |n| n[:subtree_cost] } }
+        min_tok, max_tok = group_tokens.min, group_tokens.max
+        min_cost, max_cost = group_costs.min, group_costs.max
+
+        cells = groups.each_with_index.map { |group, i|
+          primary = group.first
+          combined_tokens = group_tokens[i]
+          combined_cost = group_costs[i]
+          has_compaction = group.length > 1
+
+          color_cost = heatmap_color(combined_cost, min_cost, max_cost)
+          color_token = heatmap_color_token(combined_tokens, min_tok, max_tok)
+
+          # x/w spans all nodes in the group
+          ox = primary[:x]
+          ow = group.last[:x] + group.last[:w] - primary[:x]
+          cx = primary[:cost_x]
+          cw = group.last[:cost_x] + group.last[:cost_w] - primary[:cost_x]
+
+          num = @thread_numbers&.[](primary[:token].uuid)
+          tok_str = fmt(combined_tokens)
+          cost_str = (combined_cost > 0) ? " \u00b7 #{fmt_cost(combined_cost)}" : ""
+          compact_note = has_compaction ? " \u00b7 \u21ba compaction" : ""
+          base = num ? "Thread #{num}" : "Prompt #{i + 1}"
+          tip_text = escape_html(escape_js("#{base} \u00b7 #{tok_str} tokens#{cost_str}#{compact_note}"))
+          tip_prompt = escape_html(escape_js(primary[:token].human_text))
+          prompt_search = escape_html(truncate(primary[:token].human_text, 300).downcase)
+
+          badge = has_compaction ? %(<span class="hm-compact-badge">\u21ba</span>) : ""
+          %(<div class="hm-cell" tabindex="0" data-idx="#{i}" data-cost="#{combined_cost}" data-tokens="#{combined_tokens}" data-prompt="#{prompt_search}" data-color-cost="#{color_cost}" data-color-token="#{color_token}" data-ox="#{ox}" data-ow="#{ow}" data-cx="#{cx}" data-cw="#{cw}" data-tip="#{tip_text}" data-tip-prompt="#{tip_prompt}" style="background-color:#{color_token}" onmouseover="hmTip(this)" onmouseout="tip('')" onclick="openPrompt(#{i})"><span class="hm-idx">#{i + 1}</span>#{badge}</div>)
+        }.join("\n")
+
+        <<~HTML
+          <div id="heatmap" class="heatmap">
+            <div class="hm-meta">
+              <span class="hm-hint">Hover to preview &middot; Click to open flame graph</span>
+            </div>
+            <div class="hm-controls">
+              <span class="hm-section-label">Prompts</span>
+              <input type="text" id="hm-search" class="hm-search" placeholder="Search prompts\u2026" oninput="filterHeatmap(this.value)">
+              <button class="hm-sort-btn" id="hm-sort-btn" onclick="sortHeatmap()">&#x21C5; Sort by tokens</button>
+              <div class="hm-legend" id="hm-legend">
+                <span class="hm-legend-label" id="hm-legend-lo">few tokens</span>
+                <span class="hm-legend-ramp" id="hm-legend-ramp" data-token-grad="linear-gradient(to right, #073030, #00CED1)" data-cost-grad="linear-gradient(to right, #200510, #FF1493)" style="background:linear-gradient(to right, #073030, #00CED1)"></span>
+                <span class="hm-legend-label" id="hm-legend-hi">many</span>
+              </div>
+            </div>
+            <div class="hm-grid" id="hm-grid">
+              #{cells}
+            </div>
+            <div class="hm-empty" id="hm-empty">Click a prompt cell to explore its flame graph</div>
+          </div>
+        HTML
+      end
+
       def css
         <<~CSS
           :root {
@@ -440,6 +569,7 @@ module TokenLens
             --bar-assistant: #00CED1;
             --bar-tool: #00A8E8;
             --bar-sidechain: #C97BFF;
+            --bar-compaction: #607B8B;
             --bar-user: #1a1a2e;
             --bg: #000000;
             --surface: #0a0a14;
@@ -466,7 +596,7 @@ module TokenLens
           }
           * { box-sizing: border-box; margin: 0; padding: 0; }
           html, body { height: 100%; }
-          body { background: var(--bg); font-family: 'JetBrains Mono', 'Cascadia Mono', 'Fira Code', ui-monospace, monospace; font-size: #{FONT_SIZE}px; display: flex; flex-direction: column; min-height: 100vh; color: var(--text); }
+          body { background: var(--bg); font-family: 'JetBrains Mono', 'Cascadia Mono', 'Fira Code', ui-monospace, monospace; font-size: #{FONT_SIZE}px; display: flex; flex-direction: column; height: 100%; color: var(--text); }
           .header { display: flex; align-items: center; justify-content: space-between; padding: 4px 8px; border-bottom: 1px solid var(--border); }
           .summary { color: var(--text-dim); font-size: 11px; line-height: 20px; }
           .header-btns { display: flex; gap: 6px; align-items: center; }
@@ -477,33 +607,67 @@ module TokenLens
           .reset-btn:hover { background: var(--surface); color: var(--text); border-color: var(--accent); }
           .theme-btn { background: none; border: 1px solid var(--border); color: var(--text-dim); border-radius: 3px; padding: 2px 8px; font-size: 10px; cursor: pointer; font-family: inherit; }
           .theme-btn:hover { background: var(--surface); color: var(--text); border-color: var(--accent2); }
-          .spacer { flex: 1; }
-          .flame { position: relative; width: 100%; }
+          .flame-wrap { flex: 1; display: flex; flex-direction: column; justify-content: flex-end; overflow-x: auto; overflow-y: hidden; width: 100%; }
+          .flame { position: relative; overflow: hidden; flex-shrink: 0; }
+          .heatmap { flex-shrink: 0; padding: 12px 8px 8px; background: var(--bg); border-bottom: 1px solid var(--border); }
+          .hm-meta { display: flex; align-items: center; justify-content: space-between; margin-bottom: 6px; }
+          .hm-controls { display: flex; gap: 8px; align-items: center; justify-content: center; margin-bottom: 8px; }
+          .hm-search { background: var(--surface); border: 1px solid var(--border); color: var(--text); border-radius: 3px; padding: 3px 8px; font-size: 11px; font-family: inherit; outline: none; width: 220px; }
+          .hm-search:focus { border-color: var(--accent2); }
+          .hm-sort-btn { background: none; border: 1px solid var(--border); color: var(--text-dim); border-radius: 3px; padding: 2px 8px; font-size: 10px; cursor: pointer; font-family: inherit; }
+          .hm-sort-btn:hover { background: var(--surface); color: var(--text); border-color: var(--accent2); }
+          .hm-sort-btn.active { border-color: var(--accent2); color: var(--accent2); }
+          .hm-grid { display: flex; flex-wrap: wrap; gap: 3px; justify-content: center; }
+          .hm-cell { width: 32px; height: 32px; border-radius: 3px; cursor: pointer; transition: opacity 0.1s ease-out, transform 0.1s ease-out, box-shadow 0.1s ease-out; flex-shrink: 0; box-shadow: inset 0 0 0 1px rgba(255,255,255,0.10); display: flex; align-items: center; justify-content: center; position: relative; }
+          .hm-cell:hover { transform: scale(1.08); z-index: 10; box-shadow: inset 0 0 0 1px rgba(255,255,255,0.10), 0 0 8px var(--accent-faint); }
+          .hm-cell:focus-visible { outline: 2px solid var(--accent2); outline-offset: 1px; z-index: 10; }
+          .hm-cell.hm-active { box-shadow: 0 0 0 2px var(--accent); transform: scale(1.15); z-index: 11; }
+          .hm-idx { font-size: 9px; font-weight: 700; color: rgba(255,255,255,0.6); pointer-events: none; line-height: 1; user-select: none; }
+          .hm-compact-badge { position: absolute; top: 2px; right: 3px; font-size: 8px; color: rgba(255,255,255,0.75); pointer-events: none; line-height: 1; }
+          .hm-cell.hm-dimmed { opacity: 0.15; }
+          .heatmap.hm-strip { padding: 4px 8px; border-bottom: 1px solid var(--border); }
+          .heatmap.hm-strip .hm-meta { display: none; }
+          .heatmap.hm-strip .hm-controls { display: none; }
+          .heatmap.hm-strip .hm-grid { gap: 2px; }
+          .heatmap.hm-strip .hm-cell { width: 16px; height: 16px; border-radius: 2px; }
+          .hm-back { padding: 4px 12px; font-size: 11px; color: var(--accent); cursor: pointer; border-bottom: 1px solid var(--border); background: var(--surface); user-select: none; }
+          .hm-back:hover { background: var(--border); }
+          .heatmap.hm-strip .hm-idx { display: none; }
+          .hm-section-label { color: var(--text-dim); font-size: 10px; font-weight: 600; letter-spacing: 0.06em; text-transform: uppercase; white-space: nowrap; }
+          .hm-hint { color: var(--text-dim); font-size: 10px; opacity: 0.6; white-space: nowrap; }
+          .hm-legend { display: flex; align-items: center; gap: 4px; }
+          .hm-legend-label { color: var(--text-dim); font-size: 9px; opacity: 0.7; white-space: nowrap; }
+          .hm-legend-ramp { width: 64px; height: 8px; border-radius: 2px; flex-shrink: 0; }
+          .hm-empty { color: var(--text-dim); font-size: 12px; text-align: center; padding: 40px 0 20px; opacity: 0.5; }
+          .heatmap.hm-strip .hm-empty { display: none; }
           .bar {
             position: absolute;
-            height: #{ROW_HEIGHT - 2}px;
-            border-radius: 1px;
-            border-right: 1px solid var(--bar-border);
-            border-bottom: 2px solid var(--bar-border);
+            height: #{ROW_HEIGHT}px;
+            border-radius: 0;
+            border-right: 1.5px solid rgba(0,0,0,0.5);
+            box-shadow: inset 0 1px 0 rgba(255,255,255,0.10), inset 0 -1px 0 rgba(0,0,0,0.25);
             cursor: pointer;
             overflow: hidden;
             user-select: none;
+            transition: left 0.2s ease, width 0.2s ease;
           }
-          .bar:hover { filter: brightness(1.15) saturate(1.1); }
+          .bar:hover { filter: brightness(1.18) saturate(1.1); }
           .bar-reread { box-shadow: inset 0 -2px 0 var(--accent); }
-          .bar-compaction { box-shadow: 0 0 8px 3px rgba(255,20,147,0.7), inset 0 0 0 1px var(--accent); }
+          .bar-compaction { background: linear-gradient(to right, #7a96a8 0%, #607B8B 12%, #607B8B 88%, #4a6070 100%) !important; box-shadow: inset 0 1px 0 rgba(255,255,255,0.12), inset 0 -1px 0 rgba(0,0,0,0.3); }
           .bar-pressure { box-shadow: 0 0 5px 1px rgba(255,20,147,0.4); }
+          .bar-alt { filter: brightness(0.82); }
           .warn-badge { position: absolute; top: 1px; right: 3px; font-size: 10px; line-height: 1; color: var(--accent); pointer-events: none; font-weight: 600; }
-          .bar-c-human { background: var(--bar-human); }
-          .bar-c-task { background: var(--bar-task); }
-          .bar-c-assistant { background: var(--bar-assistant); }
-          .bar-c-tool { background: var(--bar-tool); }
-          .bar-c-sidechain { background: var(--bar-sidechain); }
-          .bar-c-user { background: var(--bar-user); }
+          .bar-c-human    { background: linear-gradient(to right, #ff3faa 0%, #FF1493 12%, #FF1493 88%, #c8107a 100%); }
+          .bar-c-task     { background: linear-gradient(to right, #ffc84e 0%, #FFB020 12%, #FFB020 88%, #d9941a 100%); }
+          .bar-c-assistant{ background: linear-gradient(to right, #1ee8eb 0%, #00CED1 12%, #00CED1 88%, #00abb0 100%); }
+          .bar-c-tool     { background: linear-gradient(to right, #22c4ff 0%, #00A8E8 12%, #00A8E8 88%, #0088c0 100%); }
+          .bar-c-sidechain{ background: linear-gradient(to right, #d88fff 0%, #C97BFF 12%, #C97BFF 88%, #a85ee0 100%); }
+          .bar-c-user     { background: var(--bar-user); }
+          .bar-c-sidechain .lbl, .bar-c-tool .lbl { color: rgba(255,255,255,0.9); }
           .lbl {
             display: block;
             padding: 0 5px;
-            line-height: #{ROW_HEIGHT - 4}px;
+            line-height: #{ROW_HEIGHT}px;
             overflow: hidden;
             text-overflow: ellipsis;
             white-space: nowrap;
@@ -525,10 +689,13 @@ module TokenLens
           .legend { padding: 2px 8px 6px; display: flex; gap: 14px; flex-wrap: wrap; }
           .legend-item { display: flex; align-items: center; gap: 5px; color: var(--text-dim); font-size: 10px; }
           .legend-swatch { display: inline-block; width: 10px; height: 10px; border-radius: 1px; flex-shrink: 0; }
+          .legend-swatch.bar-compaction { background: var(--bar-compaction) !important; }
           .floattip { position: fixed; background: var(--surface); color: var(--text); border: 1px solid var(--border); border-radius: 3px; padding: 3px 8px; font-size: 11px; font-family: inherit; pointer-events: none; display: none; white-space: nowrap; z-index: 1000; box-shadow: 0 4px 16px var(--shadow); }
           .summary-btn { background: none; border: 1px solid var(--border); color: var(--text-dim); border-radius: 3px; padding: 2px 8px; font-size: 10px; cursor: pointer; font-family: inherit; }
           .summary-btn:hover { background: var(--surface); color: var(--text); border-color: var(--accent); }
           .summary-btn.active { border-color: var(--accent); color: var(--accent); }
+          .export-btn { background: none; border: 1px solid var(--border); color: var(--text-dim); border-radius: 3px; padding: 2px 8px; font-size: 10px; cursor: pointer; font-family: inherit; }
+          .export-btn:hover { background: var(--surface); color: var(--accent2); border-color: var(--accent2); }
           .summary-panel { position: fixed; right: 0; top: 0; height: 100%; width: 260px; background: var(--surface); border-left: 1px solid var(--border); z-index: 500; overflow-y: auto; padding: 16px; box-shadow: -8px 0 24px var(--shadow); }
           .summary-panel-title { color: var(--accent); font-size: 12px; font-weight: 600; margin-bottom: 14px; display: flex; justify-content: space-between; align-items: center; letter-spacing: 0.04em; }
           .summary-close { background: none; border: none; color: var(--text-dim); font-size: 16px; cursor: pointer; padding: 0 2px; line-height: 1; font-family: inherit; }
@@ -541,12 +708,13 @@ module TokenLens
       end
 
       def js
-        w = COORD_WIDTH
-        min_pct = MIN_LABEL_PCT
+        w = @canvas_width
+        min_lbl_px = MIN_LABEL_PX
         <<~JS
           (function() {
-            var W = #{w}, MIN_PCT = #{min_pct};
-            var costMode = false;
+            var W = #{w}, MIN_LBL_PX = #{min_lbl_px};
+            var costMode = false, zoomedEl = null;
+            var hmActiveIdx = -1, hmSorted = false, hmCount = #{@hm_count};
             function bars() { return Array.from(document.querySelectorAll('.bar:not(.total-bar)')); }
             function applyBar(el, nx, nw) {
               if (nx + nw <= 0 || nx >= W) {
@@ -554,12 +722,19 @@ module TokenLens
               } else {
                 el.style.display = '';
                 el.style.left = (nx / W * 100) + '%';
-                el.style.width = Math.max(nw, 1) / W * 100 + '%';
+                el.style.width = 'calc(' + (Math.max(nw, 1) / W * 100) + '% + 1px)';
                 var lbl = el.querySelector('.lbl');
-                if (lbl) lbl.style.display = (nw / W * 100 < MIN_PCT) ? 'none' : '';
+                if (lbl) lbl.style.display = (nw < MIN_LBL_PX) ? 'none' : '';
               }
             }
             function resetBtn() { return document.getElementById('reset-btn'); }
+            function withoutTransition(fn) {
+              document.querySelectorAll('.bar').forEach(function(b) { b.style.transition = 'none'; });
+              fn();
+              var first = document.querySelector('.bar');
+              if (first) first.offsetHeight;
+              document.querySelectorAll('.bar').forEach(function(b) { b.style.transition = ''; });
+            }
             window.toggleTheme = function() {
               var root = document.documentElement;
               var isLight = root.getAttribute('data-theme') === 'light';
@@ -576,29 +751,52 @@ module TokenLens
               if (btn) btn.classList.toggle('active', !shown);
             };
             window.toggleCostView = function() {
+              var wasZoomed = zoomedEl;
+              if (wasZoomed) unzoom();
               costMode = !costMode;
-              bars().forEach(function(b) {
-                b.removeAttribute('ox');
-                b.removeAttribute('ow');
-                var nx = costMode ? +b.getAttribute('data-cx') : +b.getAttribute('data-ox');
-                var nw = costMode ? +b.getAttribute('data-cw') : +b.getAttribute('data-ow');
-                applyBar(b, nx, nw);
-                var lbl = b.querySelector('.lbl');
-                if (lbl) {
-                  var text = b.getAttribute(costMode ? 'data-cost-lbl' : 'data-token-lbl');
-                  if (text !== null) lbl.textContent = text;
-                }
+              withoutTransition(function() {
+                bars().forEach(function(b) {
+                  var nx = costMode ? +b.getAttribute('data-cx') : +b.getAttribute('data-ox');
+                  var nw = costMode ? +b.getAttribute('data-cw') : +b.getAttribute('data-ow');
+                  applyBar(b, nx, nw);
+                  var lbl = b.querySelector('.lbl');
+                  if (lbl) {
+                    var text = b.getAttribute(costMode ? 'data-cost-lbl' : 'data-token-lbl');
+                    if (text !== null) lbl.textContent = text;
+                  }
+                });
               });
               var tl = document.getElementById('total-lbl');
               if (tl) tl.innerHTML = tl.getAttribute(costMode ? 'data-cost-text' : 'data-token-text');
               var cb = document.getElementById('cost-btn');
               if (cb) { cb.textContent = costMode ? '# Token view' : '$ Cost view'; cb.classList.toggle('active', costMode); }
-              var rb = resetBtn(); if (rb) rb.style.display = 'none';
+              hmCells().forEach(function(c) { c.style.backgroundColor = c.getAttribute(costMode ? 'data-color-cost' : 'data-color-token'); });
+              var sb = document.getElementById('hm-sort-btn');
+              if (sb && hmSorted) sb.innerHTML = costMode ? '&#x21C5; Sort by cost' : '&#x21C5; Sort by tokens';
+              var ramp = document.getElementById('hm-legend-ramp');
+              if (ramp) ramp.style.background = costMode ? ramp.getAttribute('data-cost-grad') : ramp.getAttribute('data-token-grad');
+              var lmLo = document.getElementById('hm-legend-lo');
+              var lmHi = document.getElementById('hm-legend-hi');
+              if (lmLo) lmLo.textContent = costMode ? 'cheap' : 'few tokens';
+              if (lmHi) lmHi.textContent = costMode ? 'costly' : 'many';
+              scaleHeatmap();
+              if (wasZoomed) {
+                if (wasZoomed.classList && wasZoomed.classList.contains('hm-cell')) {
+                  openPrompt(+wasZoomed.getAttribute('data-idx'));
+                } else {
+                  zoom(wasZoomed);
+                }
+              }
             };
             window.zoom = function(el) {
               var fx = costMode ? +el.getAttribute('data-cx') : +el.getAttribute('data-ox');
               var fw = costMode ? +el.getAttribute('data-cw') : +el.getAttribute('data-ow');
               if (fw >= W - 1) { unzoom(); return; }
+              var wrap = document.getElementById('flame-wrap');
+              var flame = document.querySelector('.flame');
+              if (!flame.getAttribute('data-orig-w')) flame.setAttribute('data-orig-w', flame.style.width);
+              flame.style.width = wrap.clientWidth + 'px';
+              wrap.scrollLeft = 0;
               bars().forEach(function(b) {
                 if (!b.getAttribute('ox')) {
                   b.setAttribute('ox', costMode ? +b.getAttribute('data-cx') : +b.getAttribute('data-ox'));
@@ -607,8 +805,12 @@ module TokenLens
                 applyBar(b, (+b.getAttribute('ox') - fx) / fw * W, +b.getAttribute('ow') / fw * W);
               });
               var btn = resetBtn(); if (btn) btn.style.display = 'inline-block';
+              zoomedEl = el;
             };
             window.unzoom = function() {
+              var flame = document.querySelector('.flame');
+              var origW = flame.getAttribute('data-orig-w');
+              if (origW) { flame.style.width = origW; flame.removeAttribute('data-orig-w'); }
               bars().forEach(function(b) {
                 var ox = b.getAttribute('ox');
                 if (ox) {
@@ -618,12 +820,65 @@ module TokenLens
                 }
               });
               var btn = resetBtn(); if (btn) btn.style.display = 'none';
+              zoomedEl = null;
+            };
+            // resetZoom: if sub-zoomed inside a prompt → unzoom back to prompt view;
+            // if viewing a prompt (or full overview) → close back to heatmap
+            window.resetZoom = function() {
+              if (zoomedEl && zoomedEl.classList && !zoomedEl.classList.contains('hm-cell')) {
+                unzoom();
+              } else {
+                closePrompt();
+              }
+            };
+            function escSVG(s) { return s.replace(/&/g,'&amp;').replace(/</g,'&lt;').replace(/>/g,'&gt;').replace(/"/g,'&quot;'); }
+            window.exportSVG = function() {
+              var SVG_W = 1600, HDR = 28, MIN_H = 400;
+              var flame = document.querySelector('.flame');
+              var flameH = parseFloat(flame.style.height) || flame.offsetHeight;
+              var svgH = Math.max(MIN_H, Math.round(flameH) + HDR);
+              var yOffset = svgH - Math.round(flameH) - HDR; // push bars to bottom
+              var allBars = Array.from(document.querySelectorAll('.flame .bar')).filter(function(b) { return b.style.display !== 'none'; });
+              if (!allBars.length) return;
+              var st = getComputedStyle(document.documentElement);
+              var bg = st.getPropertyValue('--bg').trim() || '#000';
+              var surface = st.getPropertyValue('--surface').trim() || '#0a0a14';
+              var border = st.getPropertyValue('--border').trim() || '#1a1a2e';
+              var textDim = st.getPropertyValue('--text-dim').trim() || '#8892a0';
+              var barText = st.getPropertyValue('--bar-text').trim() || '#000';
+              var summaryEl = document.querySelector('.summary');
+              var summaryStr = summaryEl ? summaryEl.textContent.trim() : '';
+              var defs = [], rects = [], texts = [];
+              rects.push('<rect x="0" y="0" width="' + SVG_W + '" height="' + HDR + '" fill="' + surface + '"/>');
+              rects.push('<line x1="0" y1="' + HDR + '" x2="' + SVG_W + '" y2="' + HDR + '" stroke="' + border + '" stroke-width="1"/>');
+              texts.push('<text x="8" y="18" font-family="monospace" font-size="11" fill="' + textDim + '">' + escSVG(summaryStr) + '</text>');
+              allBars.forEach(function(bar, i) {
+                var x = Math.round(parseFloat(bar.style.left) / 100 * SVG_W);
+                var w = Math.max(1, Math.round(parseFloat(bar.style.width) / 100 * SVG_W));
+                var y = Math.round(parseFloat(bar.style.top)) + HDR + yOffset;
+                var h = Math.round(parseFloat(getComputedStyle(bar).height));
+                if (isNaN(x) || isNaN(y) || isNaN(w) || h <= 0) return;
+                var fill = getComputedStyle(bar).backgroundColor;
+                var cid = 'c' + i;
+                defs.push('<clipPath id="' + cid + '"><rect x="' + x + '" y="' + y + '" width="' + w + '" height="' + h + '"/></clipPath>');
+                rects.push('<rect x="' + x + '" y="' + y + '" width="' + w + '" height="' + h + '" fill="' + fill + '"/>');
+                var lbl = bar.querySelector('.lbl');
+                if (lbl && lbl.style.display !== 'none' && w > 28 && lbl.textContent.trim()) {
+                  var isTotal = bar.classList.contains('total-bar');
+                  texts.push('<text x="' + (x+3) + '" y="' + (y + h/2 + 4) + '" font-family="monospace" font-size="11" fill="' + (isTotal ? textDim : barText) + '" font-weight="' + (isTotal ? '600' : '400') + '" clip-path="url(#' + cid + ')">' + escSVG(lbl.textContent) + '</text>');
+                }
+              });
+              var svg = '<?xml version="1.0" encoding="UTF-8"?><svg xmlns="http://www.w3.org/2000/svg" width="' + SVG_W + '" height="' + svgH + '"><rect width="100%" height="100%" fill="' + bg + '"/><defs>' + defs.join('') + '</defs>' + rects.join('') + texts.join('') + '</svg>';
+              var a = document.createElement('a');
+              a.href = URL.createObjectURL(new Blob([svg], {type: 'image/svg+xml'}));
+              a.download = 'token-lens.svg';
+              document.body.appendChild(a); a.click(); document.body.removeChild(a);
             };
             function esc(t) { return t.replace(/&/g,'&amp;').replace(/</g,'&lt;').replace(/>/g,'&gt;'); }
             window.tip = function(s, prompt) {
               var el = document.getElementById('tip');
               if (!el) return;
-              if (!s && !prompt) { el.innerHTML = '&nbsp;'; return; }
+              if (!s && !prompt) { el.innerHTML = '<span class="tip-label">Hover for details \u00b7 Click to zoom</span>'; return; }
               var sep = '<span class="tip-sep">\u00b7</span>';
               var parts = s ? s.split(' | ').filter(Boolean) : [];
               var html = parts.map(function(p, i) {
@@ -634,7 +889,7 @@ module TokenLens
                 if (/^[^:]*:\s/.test(p)) return '<span class="tip-label">' + e + '</span>';
                 return '<span class="tip-code">' + e + '</span>';
               }).join(sep);
-              if (prompt) html += '<span class="tip-prompt">' + esc(prompt) + '</span>';
+              if (prompt) html += (html ? sep : '') + '<span class="tip-prompt">' + esc(prompt) + '</span>';
               el.innerHTML = html;
             };
             var mx = 0, my = 0;
@@ -659,6 +914,134 @@ module TokenLens
               var bar = e.target.closest && e.target.closest('.bar:not(.total-bar)');
               if (bar) { var ft = document.getElementById('ftip'); if (ft) ft.style.display = 'none'; }
             });
+            // Disable transitions on initial load
+            document.querySelectorAll('.bar').forEach(function(b) { b.style.transition = 'none'; });
+            requestAnimationFrame(function() { requestAnimationFrame(function() {
+              document.querySelectorAll('.bar').forEach(function(b) { b.style.transition = ''; });
+            }); });
+            // Heatmap
+            function hmCells() { return Array.from(document.querySelectorAll('.hm-cell')); }
+            window.hmTip = function(el) { tip(el.getAttribute('data-tip'), el.getAttribute('data-tip-prompt')); };
+            function zoomToRoot(cell) {
+              var fx = costMode ? +cell.getAttribute('data-cx') : +cell.getAttribute('data-ox');
+              var fw = costMode ? +cell.getAttribute('data-cw') : +cell.getAttribute('data-ow');
+              if (fw <= 0) return;
+              var wrap = document.getElementById('flame-wrap');
+              var flame = document.querySelector('.flame');
+              if (!flame.getAttribute('data-orig-w')) flame.setAttribute('data-orig-w', flame.style.width);
+              flame.style.width = wrap.clientWidth + 'px';
+              wrap.scrollLeft = 0;
+              bars().forEach(function(b) {
+                if (!b.getAttribute('ox')) {
+                  b.setAttribute('ox', costMode ? +b.getAttribute('data-cx') : +b.getAttribute('data-ox'));
+                  b.setAttribute('ow', costMode ? +b.getAttribute('data-cw') : +b.getAttribute('data-ow'));
+                }
+                applyBar(b, (+b.getAttribute('ox') - fx) / fw * W, +b.getAttribute('ow') / fw * W);
+              });
+              zoomedEl = cell;
+            }
+            window.openPrompt = function(idx) {
+              hmActiveIdx = idx;
+              var hm = document.getElementById('heatmap');
+              if (hm) hm.classList.add('hm-strip');
+              hmCells().forEach(function(c) { c.classList.toggle('hm-active', +c.getAttribute('data-idx') === idx); });
+              var fw = document.getElementById('flame-wrap');
+              if (fw) fw.style.display = '';
+              var lg = document.getElementById('legend'); if (lg) lg.style.display = '';
+              var back = document.getElementById('hm-back');
+              if (back) back.style.display = '';
+              var cell = document.querySelector('.hm-cell[data-idx="' + idx + '"]');
+              if (cell) { withoutTransition(function() { zoomToRoot(cell); }); }
+              var btn = resetBtn(); if (btn) btn.style.display = 'none';
+            };
+            window.closePrompt = function() {
+              hmActiveIdx = -1;
+              unzoom();
+              var fw = document.getElementById('flame-wrap');
+              if (fw) fw.style.display = 'none';
+              var lg = document.getElementById('legend'); if (lg) lg.style.display = 'none';
+              var hm = document.getElementById('heatmap');
+              if (hm) hm.classList.remove('hm-strip');
+              hmCells().forEach(function(c) { c.classList.remove('hm-active'); });
+              var back = document.getElementById('hm-back');
+              if (back) back.style.display = 'none';
+              var si = document.getElementById('hm-search');
+              if (si) { si.value = ''; filterHeatmap(''); }
+            };
+            window.sortHeatmap = function() {
+              hmSorted = !hmSorted;
+              var grid = document.getElementById('hm-grid');
+              if (!grid) return;
+              var cells = hmCells();
+              if (hmSorted) {
+                cells.sort(function(a, b) {
+                  var av = costMode ? parseFloat(a.getAttribute('data-cost')) : parseInt(a.getAttribute('data-tokens'));
+                  var bv = costMode ? parseFloat(b.getAttribute('data-cost')) : parseInt(b.getAttribute('data-tokens'));
+                  return bv - av;
+                });
+              } else {
+                cells.sort(function(a, b) { return parseInt(a.getAttribute('data-idx')) - parseInt(b.getAttribute('data-idx')); });
+              }
+              cells.forEach(function(c) { grid.appendChild(c); });
+              scaleHeatmap();
+              var btn = document.getElementById('hm-sort-btn');
+              if (btn) btn.classList.toggle('active', hmSorted);
+            };
+            window.filterHeatmap = function(query) {
+              var q = query.toLowerCase().trim();
+              hmCells().forEach(function(c) {
+                var match = !q || (c.getAttribute('data-prompt') || '').indexOf(q) !== -1;
+                c.classList.toggle('hm-dimmed', !match);
+              });
+              scaleHeatmap();
+            };
+            function scaleHeatmap() {
+              var hm = document.getElementById('heatmap');
+              if (hm && hm.classList.contains('hm-strip')) return;
+              var allCells = hmCells();
+              var cells = allCells.filter(function(c) { return !c.classList.contains('hm-dimmed'); });
+              var n = cells.length || 1;
+              var maxCell = Math.min(96, Math.max(36, Math.floor(1400 / Math.sqrt(n))));
+              var minCell = 16;
+              var vals = cells.map(function(c) {
+                return costMode ? parseFloat(c.getAttribute('data-cost')) : parseInt(c.getAttribute('data-tokens'));
+              });
+              var maxVal = Math.max.apply(null, vals) || 1;
+              var minVal = Math.min.apply(null, vals) || 0;
+              cells.forEach(function(c, i) {
+                var v = vals[i];
+                var t = (maxVal === minVal) ? 0.5 : (v - minVal) / (maxVal - minVal);
+                var sz = Math.round(minCell + (maxCell - minCell) * Math.sqrt(t));
+                c.style.width = sz + 'px';
+                c.style.height = sz + 'px';
+                var lbl = c.querySelector('.hm-idx');
+                if (lbl) lbl.style.display = (sz < 24) ? 'none' : '';
+              });
+              allCells.filter(function(c) { return c.classList.contains('hm-dimmed'); }).forEach(function(c) {
+                c.style.width = minCell + 'px';
+                c.style.height = minCell + 'px';
+              });
+            }
+            document.addEventListener('keydown', function(e) {
+              if (e.key === 'Enter' && document.activeElement && document.activeElement.classList.contains('hm-cell')) {
+                e.preventDefault();
+                openPrompt(+document.activeElement.getAttribute('data-idx'));
+                return;
+              }
+              if (hmActiveIdx < 0) return;
+              if (e.target && e.target.tagName === 'INPUT') return;
+              if (e.key === 'ArrowLeft' || e.key === 'ArrowUp') {
+                e.preventDefault();
+                if (hmActiveIdx > 0) openPrompt(hmActiveIdx - 1);
+              } else if (e.key === 'ArrowRight' || e.key === 'ArrowDown') {
+                e.preventDefault();
+                if (hmActiveIdx < hmCount - 1) openPrompt(hmActiveIdx + 1);
+              } else if (e.key === 'Escape') {
+                e.preventDefault();
+                closePrompt();
+              }
+            });
+            scaleHeatmap();
           })();
         JS
       end

data/lib/token_lens/renderer/layout.rb

@@ -5,6 +5,7 @@ module TokenLens
     class Layout
       CANVAS_WIDTH = 1200
       ROW_HEIGHT = 32
+      MIN_THREAD_WIDTH = 80  # minimum pixels per root thread so slivers are readable
 
       def initialize(canvas_width: CANVAS_WIDTH)
         @canvas_width = canvas_width
@@ -12,15 +13,17 @@ module TokenLens
 
       def layout(nodes)
         max_depth = all_nodes(nodes).map { |n| n[:depth] }.max || 0
+        effective_width = [@canvas_width, nodes.length * MIN_THREAD_WIDTH].max
+
         total = nodes.sum { |n| n[:subtree_tokens] }
-        scale = (total > 0) ? @canvas_width.to_f / total : 1.0
+        scale = (total > 0) ? effective_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
+        cost_scale = (total_cost > 0) ? effective_width.to_f / total_cost : 1.0
         position_cost(nodes, x: 0, scale: cost_scale, max_depth: max_depth)
 
-        nodes
+        effective_width
       end
 
       private
@@ -28,23 +31,25 @@ module TokenLens
       # 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
+        cursor = x.to_f
         nodes.each do |node|
-          node[:x] = cursor
+          start = cursor.round
+          cursor += node[:subtree_tokens] * scale
+          node[:x] = start
           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]
+          node[:w] = cursor.round - start
+          position(node[:children], x: node[:x], scale: scale, max_depth: max_depth)
         end
       end
 
       def position_cost(nodes, x:, scale:, max_depth:)
-        cursor = x
+        cursor = x.to_f
         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]
+          start = cursor.round
+          cursor += node[:subtree_cost] * scale
+          node[:cost_x] = start
+          node[:cost_w] = cursor.round - start
+          position_cost(node[:children], x: node[:cost_x], scale: scale, max_depth: max_depth)
         end
       end
 

data/lib/token_lens/renderer/reshaper.rb

@@ -5,7 +5,17 @@ module TokenLens
     class Reshaper
       def reshape(nodes)
         nodes = collapse_streaming(nodes)
-        nodes.flat_map { |node| process_root(node) }
+        # Process roots iteratively so that human prompts discovered mid-thread
+        # (nested inside an assistant chain) are hoisted to the top level rather
+        # than stacked as children of the prompt that preceded them.
+        @pending_roots = nodes.dup
+        result = []
+        while @pending_roots.any?
+          batch = @pending_roots
+          @pending_roots = []
+          result += batch.flat_map { |node| process_root(node) }
+        end
+        result
       end
 
       private
@@ -58,11 +68,17 @@ module TokenLens
       # 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:)
+      #
+      # through_assistant: tracks whether we've passed at least one assistant turn.
+      # A human prompt encountered BEFORE any assistant (e.g. a screenshot attached
+      # to the same user turn) is treated as transparent — we recurse into its
+      # children rather than hoisting it. A human prompt encountered AFTER an
+      # assistant is a genuine new conversational turn and gets hoisted.
+      def flatten_thread(nodes, prev_input:, through_assistant: false)
         nodes.flat_map do |node|
           t = node[:token]
           if t.role == "user" && !t.is_human_prompt?
-            flatten_thread(node[:children], prev_input: prev_input)
+            flatten_thread(node[:children], prev_input: prev_input, through_assistant: through_assistant)
           elsif t.role == "assistant"
             marginal = [t.input_tokens - prev_input, 0].max
             compaction = prev_input > 0 && t.input_tokens < prev_input * 0.5
@@ -77,9 +93,15 @@ module TokenLens
               token: t.with(marginal_input_tokens: marginal, is_compaction: compaction),
               children: sidechain
             )
-            [updated] + flatten_thread(chain, prev_input: t.input_tokens)
+            [updated] + flatten_thread(chain, prev_input: t.input_tokens, through_assistant: true)
+          elsif through_assistant
+            # Human prompt after an assistant — genuine new turn, hoist to top level.
+            @pending_roots << node
+            []
           else
-            process_root(node)
+            # Human prompt before any assistant (consecutive user messages, e.g. an
+            # image attachment). Treat as transparent and continue into its children.
+            flatten_thread(node[:children], prev_input: prev_input, through_assistant: false)
           end
         end
       end

data/lib/token_lens/version.rb

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

metadata

@@ -1,13 +1,13 @@
 --- !ruby/object:Gem::Specification
 name: token-lens
 version: !ruby/object:Gem::Version
-  version: 0.3.0
+  version: 0.5.0
 platform: ruby
 authors:
 - rob durst
 bindir: bin
 cert_chain: []
-date: 2026-03-23 00:00:00.000000000 Z
+date: 2026-03-24 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: thor