token-lens 0.1.0 → 0.4.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: fc58978bdc3f4759309ded99d65da22be12a01287ea79acdb43ab873e7cb8771
-  data.tar.gz: 34c54363d13e72ed669e3f27bbf1fc8af2d3665e958d977baabfaac9ab80f849
+  metadata.gz: 4e268e0550d46641067f1f99507b60d92251a4b58d025b39b8b3f90c428eb47a
+  data.tar.gz: 03243e21fd60fa99443b35a9e384fbb84d1fcf7c00072c297fa5fb0d46ae83a7
 SHA512:
-  metadata.gz: 3f3c35dedc0cd0efd6e12b8d136aab7f4328ba4619f912f07c94727e3bd6c4c2bb4184d07a9e3bbb117dd0eaeec28edd27c5f7100f140a79b12043773a066881
-  data.tar.gz: d0767d26c77b4ba6dfec6f3a40b70c7d1dde3b3a58094fa8707b4bf57bc860612ed1af91c161b31fc3f3faa78f3db25ed6f5ddc3ae2f9422ca351084e3d5ec41
+  metadata.gz: 88af830c547a6ef86aa5b0fdf763822df9a9ca134b5d304769ba66f97585c731012c90d9498a599bac64abdace620fbec9f113b9c9e476acf301c4a18dfe7bc0
+  data.tar.gz: d4c8be308c772f5c117d79a5b82d4939b5f2f45776be127f3ec6bc5bb73487055be5c29b92c30b12e848fad89f6e4f9b387ab648ae055a27e282639b66fec42a

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
@@ -10,9 +17,52 @@ Record Data --> Interpret Data --> Render Data
 
 ## Data Sources
 
-1. Tap OTEL console output
-2. Tail session JSONL
-3. Local API proxy
+1. [unsupported] Tap OTEL console output
+2. [supported] Tail session JSONL
+3. [unsupported] Local API proxy
+
+## Quick Start
+
+### Zero setup — render any session right now
+
+No recording needed. token-lens reads directly from Claude Code's session files:
+
+```
+gem install token-lens
+token-lens render
+open flame.html
+```
+
+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):
+
+```
+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
+```
 
 ## Contributing
 

data/lib/token_lens/cli.rb

@@ -6,14 +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]).run
+      Commands::Record.new(
+        duration_in_seconds: options[:duration_in_seconds],
+        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,13 +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:)
+      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
@@ -16,7 +21,7 @@ module TokenLens
         queue = Queue.new
         events = []
 
-        thread = Thread.new { Sources::Jsonl.new(queue).start }
+        thread = Thread.new { Sources::Jsonl.new(queue, project_dir: @project_dir).start }
         drain_thread = Thread.new { loop { events << queue.pop } }
 
         trap("INT") { finish(thread, drain_thread, queue, events) }
@@ -33,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,17 +5,20 @@ 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)
         Renderer::Annotator.new.annotate(tree)
         Renderer::Layout.new.layout(tree)
@@ -23,6 +26,17 @@ module TokenLens
         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.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/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/session.rb

@@ -16,15 +16,28 @@ module TokenLens
 
     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
 
+    # Returns the most recently modified session file across ALL projects.
+    def self.latest_jsonl
+      all = CLAUDE_DIR.glob("*/*.jsonl")
+      raise "No session files found in #{CLAUDE_DIR}" if all.empty?
+      all.max_by(&:mtime)
+    end
+
+    # Like active_jsonl but falls back to latest_jsonl with a warning when
+    # the current directory has no sessions (e.g. running from ~/Desktop).
+    def self.active_or_latest_jsonl(dir = Dir.pwd)
+      active_jsonl(dir)
+    rescue RuntimeError
+      path = latest_jsonl
+      warn "  [session] no sessions for #{dir}, using most recent: #{path}"
+      path
+    end
+
     def self.tail(path, &block)
       last_pos = File.size(path)
       loop do

data/lib/token_lens/sources/jsonl.rb

@@ -6,9 +6,9 @@ require "token_lens/session"
 module TokenLens
   module Sources
     class Jsonl
-      def initialize(queue)
+      def initialize(queue, project_dir: nil)
         @queue = queue
-        @path = Session.active_jsonl
+        @path = project_dir ? Session.active_jsonl(project_dir) : Session.active_or_latest_jsonl
       end
 
       def start

data/lib/token_lens/version.rb

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

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: token-lens
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.4.0
 platform: ruby
 authors:
 - rob durst