llm_cost_tracker 0.2.0.alpha2 → 0.3.0

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -18,22 +18,39 @@ module LlmCostTracker
 
         request_url  = request_env.url.to_s
         request_body = read_body(request_env.body) || ""
+        parser       = Parsers::Registry.find_for(request_url)
+        streaming    = parser&.streaming_request?(request_url, request_body)
+        stream_buffer = install_stream_tap(request_env) if streaming
 
-        enforce_budget_before_request(request_url)
+        Tracker.enforce_budget! if parser
         started_at = monotonic_time
 
         @app.call(request_env).on_complete do |response_env|
-          process(request_env, request_url, request_body, response_env, elapsed_ms(started_at))
+          process(
+            parser: parser,
+            request_env: request_env,
+            request_url: request_url,
+            request_body: request_body,
+            response_env: response_env,
+            latency_ms: elapsed_ms(started_at),
+            streaming: streaming,
+            stream_buffer: stream_buffer
+          )
         end
       end
 
       private
 
-      def process(request_env, request_url, request_body, response_env, latency_ms)
-        parser = Parsers::Registry.find_for(request_url)
+      def process(parser:, request_env:, request_url:, request_body:, response_env:,
+                  latency_ms:, streaming:, stream_buffer:)
         return unless parser
 
-        parsed = parse_response(parser, request_url, request_body, response_env)
+        parsed =
+          if streaming
+            parse_stream(parser, request_url, request_body, response_env, stream_buffer)
+          else
+            parse_response(parser, request_url, request_body, response_env)
+          end
         return unless parsed
 
         Tracker.record(
@@ -42,6 +59,8 @@ module LlmCostTracker
           input_tokens: parsed.input_tokens,
           output_tokens: parsed.output_tokens,
           latency_ms: latency_ms,
+          stream: parsed.stream,
+          usage_source: parsed.usage_source,
           metadata: resolved_tags(request_env).merge(parsed.metadata)
         )
       rescue LlmCostTracker::Error
@@ -54,7 +73,9 @@ module LlmCostTracker
         response_body = read_body(response_env.body)
         unless response_body
           Logging.warn(
-            "Unable to read response body for #{request_url}; streaming/SSE responses require manual tracking."
+            "Unable to read response body for #{request_url}; " \
+            "streaming responses are captured automatically for OpenAI/Anthropic/Gemini " \
+            "or via LlmCostTracker.track_stream for custom clients."
           )
           return nil
         end
@@ -62,10 +83,37 @@ module LlmCostTracker
         parser.parse(request_url, request_body, response_env.status, response_body)
       end
 
-      def enforce_budget_before_request(request_url)
-        return unless Parsers::Registry.find_for(request_url)
+      def parse_stream(parser, request_url, request_body, response_env, stream_buffer)
+        body = stream_buffer&.string
+        body = read_body(response_env.body) if body.nil? || body.empty?
+
+        if body.nil? || body.empty?
+          Logging.warn(
+            "Unable to capture streaming response for #{request_url}; " \
+            "fall back to LlmCostTracker.track_stream for manual capture."
+          )
+          return nil
+        end
+
+        events = Parsers::SSE.parse(body)
+        parser.parse_stream(request_url, request_body, response_env.status, events)
+      end
+
+      def install_stream_tap(request_env)
+        return nil unless request_env.respond_to?(:request) && request_env.request
 
-        Tracker.enforce_budget!
+        original = request_env.request.on_data
+        return nil unless original
+
+        buffer = StringIO.new
+        request_env.request.on_data = proc do |chunk, size, env|
+          buffer << chunk.to_s
+          original.call(chunk, size, env)
+        end
+        buffer
+      rescue StandardError => e
+        Logging.warn("Unable to install streaming tap: #{e.class}: #{e.message}")
+        nil
       end
 
       def read_body(body)

data/lib/llm_cost_tracker/parsed_usage.rb

@@ -10,11 +10,13 @@ module LlmCostTracker
     :cached_input_tokens,
     :cache_read_input_tokens,
     :cache_creation_input_tokens,
-    :reasoning_tokens
+    :reasoning_tokens,
+    :stream,
+    :usage_source
   )
 
   class ParsedUsage
-    TRACKING_KEYS = %i[provider model input_tokens output_tokens total_tokens].freeze
+    TRACKING_KEYS = %i[provider model input_tokens output_tokens total_tokens stream usage_source].freeze
 
     def self.build(**attributes)
       new(
@@ -26,7 +28,9 @@ module LlmCostTracker
         cached_input_tokens: attributes[:cached_input_tokens],
         cache_read_input_tokens: attributes[:cache_read_input_tokens],
         cache_creation_input_tokens: attributes[:cache_creation_input_tokens],
-        reasoning_tokens: attributes[:reasoning_tokens]
+        reasoning_tokens: attributes[:reasoning_tokens],
+        stream: attributes[:stream] || false,
+        usage_source: attributes[:usage_source]
       )
     end
 

data/lib/llm_cost_tracker/parsers/anthropic.rb

@@ -16,6 +16,10 @@ module LlmCostTracker
         false
       end
 
+      def provider_names
+        %w[anthropic]
+      end
+
       def parse(_request_url, request_body, response_status, response_body)
         return nil unless response_status == 200
 
@@ -33,9 +37,83 @@ module LlmCostTracker
           total_tokens: usage["input_tokens"].to_i + usage["output_tokens"].to_i +
             usage["cache_read_input_tokens"].to_i + usage["cache_creation_input_tokens"].to_i,
           cache_read_input_tokens: usage["cache_read_input_tokens"],
-          cache_creation_input_tokens: usage["cache_creation_input_tokens"]
+          cache_creation_input_tokens: usage["cache_creation_input_tokens"],
+          usage_source: :response
         )
       end
+
+      def parse_stream(_request_url, request_body, response_status, events)
+        return nil unless response_status == 200
+
+        request = safe_json_parse(request_body)
+        model = stream_model(events) || request["model"]
+        usage = stream_usage(events)
+
+        if usage
+          input = usage["input_tokens"].to_i
+          output = usage["output_tokens"].to_i
+          cache_read = usage["cache_read_input_tokens"].to_i
+          cache_creation = usage["cache_creation_input_tokens"].to_i
+
+          ParsedUsage.build(
+            provider: "anthropic",
+            model: model,
+            input_tokens: input,
+            output_tokens: output,
+            total_tokens: input + output + cache_read + cache_creation,
+            cache_read_input_tokens: usage["cache_read_input_tokens"],
+            cache_creation_input_tokens: usage["cache_creation_input_tokens"],
+            stream: true,
+            usage_source: :stream_final
+          )
+        else
+          ParsedUsage.build(
+            provider: "anthropic",
+            model: model,
+            input_tokens: 0,
+            output_tokens: 0,
+            total_tokens: 0,
+            stream: true,
+            usage_source: :unknown
+          )
+        end
+      end
+
+      private
+
+      def stream_usage(events)
+        start_usage = nil
+        latest_delta = nil
+
+        events.each do |event|
+          data = event[:data]
+          next unless data.is_a?(Hash)
+
+          case data["type"]
+          when "message_start"
+            start_usage = data.dig("message", "usage")
+          when "message_delta"
+            latest_delta = data["usage"] if data["usage"].is_a?(Hash)
+          end
+        end
+
+        return nil unless start_usage || latest_delta
+
+        (start_usage || {}).merge(latest_delta || {}) do |_key, start_val, delta_val|
+          delta_val.nil? ? start_val : delta_val
+        end
+      end
+
+      def stream_model(events)
+        events.each do |event|
+          data = event[:data]
+          next unless data.is_a?(Hash)
+
+          model = data.dig("message", "model")
+          return model if model && !model.empty?
+        end
+        nil
+      end
     end
   end
 end

data/lib/llm_cost_tracker/parsers/base.rb

@@ -5,19 +5,31 @@ require "json"
 module LlmCostTracker
   module Parsers
     class Base
-      # Parse a provider response into a {LlmCostTracker::ParsedUsage}, or return
-      # nil when the response is not trackable (non-200, missing usage, etc).
-      #
-      # @return [LlmCostTracker::ParsedUsage, nil]
       def parse(request_url, request_body, response_status, response_body)
         raise NotImplementedError
       end
 
-      # Returns true if this parser can handle the given URL.
+      def provider_names
+        []
+      end
+
       def match?(url)
         raise NotImplementedError
       end
 
+      def streaming_request?(_request_url, request_body)
+        return false if request_body.nil?
+
+        body = request_body.to_s
+        return false if body.empty?
+
+        body.include?('"stream":true') || body.include?('"stream": true') || body.include?("stream: true")
+      end
+
+      def parse_stream(_request_url, _request_body, _response_status, _events)
+        nil
+      end
+
       private
 
       def safe_json_parse(body)

data/lib/llm_cost_tracker/parsers/gemini.rb

@@ -9,6 +9,7 @@ module LlmCostTracker
     class Gemini < Base
       HOSTS = %w[generativelanguage.googleapis.com].freeze
       TRACKED_PATH_PATTERN = %r{/models/[^/:]+:(?:generateContent|streamGenerateContent)\z}
+      STREAM_PATH_PATTERN  = /:streamGenerateContent\z/
 
       def match?(url)
         uri = URI.parse(url.to_s)
@@ -17,6 +18,16 @@ module LlmCostTracker
         false
       end
 
+      def provider_names
+        %w[gemini]
+      end
+
+      def streaming_request?(request_url, request_body)
+        return true if streaming_url?(request_url)
+
+        super
+      end
+
       def parse(request_url, _request_body, response_status, response_body)
         return nil unless response_status == 200
 
@@ -24,31 +35,73 @@ module LlmCostTracker
         usage    = response["usageMetadata"]
         return nil unless usage
 
-        # Extract model from URL: /v1beta/models/gemini-2.5-flash:generateContent
+        build_parsed_usage(request_url, usage, usage_source: :response)
+      end
+
+      def parse_stream(request_url, _request_body, response_status, events)
+        return nil unless response_status == 200
+
+        usage = merged_stream_usage(events)
         model = extract_model_from_url(request_url)
 
+        if usage
+          build_parsed_usage(request_url, usage, stream: true, usage_source: :stream_final)
+        else
+          ParsedUsage.build(
+            provider: "gemini",
+            model: model,
+            input_tokens: 0,
+            output_tokens: 0,
+            total_tokens: 0,
+            stream: true,
+            usage_source: :unknown
+          )
+        end
+      end
+
+      private
+
+      def build_parsed_usage(request_url, usage, usage_source:, stream: false)
         ParsedUsage.build(
           provider: "gemini",
-          model: model,
+          model: extract_model_from_url(request_url),
           input_tokens: usage["promptTokenCount"].to_i,
           output_tokens: output_tokens(usage),
           total_tokens: usage["totalTokenCount"].to_i,
-          cached_input_tokens: usage["cachedContentTokenCount"]
+          cached_input_tokens: usage["cachedContentTokenCount"],
+          stream: stream,
+          usage_source: usage_source
         )
       end
 
-      private
+      def merged_stream_usage(events)
+        latest = nil
+        events.each do |event|
+          data = event[:data]
+          next unless data.is_a?(Hash)
+
+          meta = data["usageMetadata"]
+          latest = meta if meta.is_a?(Hash)
+        end
+        latest
+      end
 
       def output_tokens(usage)
         usage["candidatesTokenCount"].to_i + usage["thoughtsTokenCount"].to_i
       end
 
+      def streaming_url?(request_url)
+        URI.parse(request_url.to_s).path.match?(STREAM_PATH_PATTERN)
+      rescue URI::InvalidURIError
+        false
+      end
+
       def extract_model_from_url(url)
         uri = URI.parse(url.to_s)
         match = uri.path.match(%r{/models/([^/:]+)})
-        match ? match[1] : "unknown"
+        match && match[1]
       rescue URI::InvalidURIError
-        "unknown"
+        nil
       end
     end
   end

data/lib/llm_cost_tracker/parsers/openai.rb

@@ -20,10 +20,18 @@ module LlmCostTracker
         false
       end
 
+      def provider_names
+        %w[openai]
+      end
+
       def parse(request_url, request_body, response_status, response_body)
         parse_openai_usage(request_url, request_body, response_status, response_body)
       end
 
+      def parse_stream(request_url, request_body, response_status, events)
+        parse_openai_stream_usage(request_url, request_body, response_status, events)
+      end
+
       private
 
       def provider_for(_request_url)

data/lib/llm_cost_tracker/parsers/openai_compatible.rb

@@ -19,10 +19,18 @@ module LlmCostTracker
         false
       end
 
+      def provider_names
+        ["openai_compatible", *configured_providers.each_value.map(&:to_s)].uniq.freeze
+      end
+
       def parse(request_url, request_body, response_status, response_body)
         parse_openai_usage(request_url, request_body, response_status, response_body)
       end
 
+      def parse_stream(request_url, request_body, response_status, events)
+        parse_openai_stream_usage(request_url, request_body, response_status, events)
+      end
+
       private
 
       def provider_for(request_url)

data/lib/llm_cost_tracker/parsers/openai_usage.rb

@@ -20,10 +20,64 @@ module LlmCostTracker
           input_tokens: (usage["prompt_tokens"] || usage["input_tokens"]).to_i,
           output_tokens: (usage["completion_tokens"] || usage["output_tokens"]).to_i,
           total_tokens: usage["total_tokens"].to_i,
-          cached_input_tokens: cached_input_tokens(usage)
+          cached_input_tokens: cached_input_tokens(usage),
+          usage_source: :response
         )
       end
 
+      def parse_openai_stream_usage(request_url, request_body, response_status, events)
+        return nil unless response_status == 200
+
+        request = safe_json_parse(request_body)
+        model = detect_stream_model(events) || request["model"]
+        usage = detect_stream_usage(events)
+
+        if usage
+          ParsedUsage.build(
+            provider: provider_for(request_url),
+            model: model,
+            input_tokens: (usage["prompt_tokens"] || usage["input_tokens"]).to_i,
+            output_tokens: (usage["completion_tokens"] || usage["output_tokens"]).to_i,
+            total_tokens: usage["total_tokens"].to_i,
+            cached_input_tokens: cached_input_tokens(usage),
+            stream: true,
+            usage_source: :stream_final
+          )
+        else
+          ParsedUsage.build(
+            provider: provider_for(request_url),
+            model: model,
+            input_tokens: 0,
+            output_tokens: 0,
+            total_tokens: 0,
+            stream: true,
+            usage_source: :unknown
+          )
+        end
+      end
+
+      def detect_stream_usage(events)
+        events.reverse_each do |event|
+          data = event[:data]
+          next unless data.is_a?(Hash)
+
+          usage = data["usage"]
+          return usage if usage.is_a?(Hash) && !usage.empty?
+        end
+        nil
+      end
+
+      def detect_stream_model(events)
+        events.each do |event|
+          data = event[:data]
+          next unless data.is_a?(Hash)
+
+          model = data["model"]
+          return model if model && !model.to_s.empty?
+        end
+        nil
+      end
+
       def cached_input_tokens(usage)
         details = usage["prompt_tokens_details"] || usage["input_tokens_details"] || {}
         details["cached_tokens"]

data/lib/llm_cost_tracker/parsers/registry.rb

@@ -1,23 +1,35 @@
 # frozen_string_literal: true
 
+require "monitor"
+
 module LlmCostTracker
   module Parsers
     class Registry
+      MUTEX = Monitor.new
+
       class << self
         def parsers
-          @parsers ||= default_parsers
+          @parsers || MUTEX.synchronize { @parsers ||= default_parsers.freeze }
         end
 
         def register(parser)
-          parsers.unshift(parser)
+          MUTEX.synchronize do
+            current = @parsers || default_parsers.freeze
+            @parsers = ([parser] + current).freeze
+          end
         end
 
         def find_for(url)
           parsers.find { |parser| parser.match?(url) }
         end
 
+        def find_for_provider(provider)
+          provider_name = provider.to_s
+          parsers.find { |parser| parser.provider_names.include?(provider_name) }
+        end
+
         def reset!
-          @parsers = nil
+          MUTEX.synchronize { @parsers = nil }
         end
 
         private

data/lib/llm_cost_tracker/parsers/sse.rb

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+require "json"
+
+module LlmCostTracker
+  module Parsers
+    module SSE
+      DONE_MARKER = "[DONE]"
+
+      class << self
+        def parse(body)
+          return [] if body.nil? || body.empty?
+
+          return parse_json_array(body) if probably_json_array?(body)
+
+          parse_event_stream(body)
+        end
+
+        private
+
+        def parse_event_stream(body)
+          events = []
+          current_event = nil
+          data_lines = []
+
+          body.each_line do |raw|
+            line = raw.chomp
+
+            if line.empty?
+              events << finalize_event(current_event, data_lines) if data_lines.any?
+              current_event = nil
+              data_lines = []
+              next
+            end
+
+            next if line.start_with?(":")
+
+            field, _, value = line.partition(":")
+            value = value[1..] if value.start_with?(" ")
+
+            case field
+            when "event" then current_event = value
+            when "data"  then data_lines << value
+            end
+          end
+
+          events << finalize_event(current_event, data_lines) if data_lines.any?
+          events.compact
+        end
+
+        def parse_json_array(body)
+          parsed = JSON.parse(body)
+          return [] unless parsed.is_a?(Array)
+
+          parsed.map { |entry| { event: nil, data: entry } }
+        rescue JSON::ParserError
+          []
+        end
+
+        def finalize_event(event_name, data_lines)
+          payload = data_lines.join("\n")
+          return nil if payload == DONE_MARKER
+
+          { event: event_name, data: decode_data(payload) }
+        end
+
+        def decode_data(payload)
+          return payload if payload.empty?
+
+          JSON.parse(payload)
+        rescue JSON::ParserError
+          payload
+        end
+
+        def probably_json_array?(body)
+          body.lstrip.start_with?("[")
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/price_registry.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "json"
+require "monitor"
 require "yaml"
 
 require_relative "logging"
@@ -10,15 +11,18 @@ module LlmCostTracker
     DEFAULT_PRICES_PATH = File.expand_path("prices.json", __dir__)
     EMPTY_PRICES = {}.freeze
     PRICE_KEYS = %w[input cached_input output cache_read_input cache_creation_input].freeze
-    METADATA_KEYS = %w[_source _updated _notes].freeze
+    METADATA_KEYS = %w[_source _source_version _fetched_at _updated _notes _validator_override].freeze
+    MUTEX = Monitor.new
 
     class << self
       def builtin_prices
-        @builtin_prices ||= normalize_price_table(raw_registry.fetch("models", {})).freeze
+        @builtin_prices ||= MUTEX.synchronize do
+          @builtin_prices || normalize_price_table(raw_registry.fetch("models", {})).freeze
+        end
       end
 
       def metadata
-        @metadata ||= raw_registry.fetch("metadata", {}).freeze
+        @metadata ||= MUTEX.synchronize { @metadata || raw_registry.fetch("metadata", {}).freeze }
       end
 
       def normalize_price_table(table)
@@ -35,9 +39,14 @@ module LlmCostTracker
         cached = @file_prices_cache
         return cached[:value] if cached && cached[:key] == cache_key
 
-        value = normalize_file_prices(price_file_models(load_price_file(path)), path: path).freeze
-        @file_prices_cache = { key: cache_key, value: value }.freeze
-        value
+        MUTEX.synchronize do
+          cached = @file_prices_cache
+          return cached[:value] if cached && cached[:key] == cache_key
+
+          value = normalize_file_prices(price_file_models(load_price_file(path)), path: path).freeze
+          @file_prices_cache = { key: cache_key, value: value }.freeze
+          value
+        end
       rescue Errno::ENOENT, JSON::ParserError, Psych::Exception, ArgumentError, TypeError, NoMethodError => e
         raise Error, "Unable to load prices_file #{path.inspect}: #{e.message}"
       end
@@ -45,7 +54,9 @@ module LlmCostTracker
       private
 
       def raw_registry
-        @raw_registry ||= JSON.parse(File.read(DEFAULT_PRICES_PATH)).freeze
+        @raw_registry ||= MUTEX.synchronize do
+          @raw_registry || JSON.parse(File.read(DEFAULT_PRICES_PATH)).freeze
+        end
       end
 
       def normalize_price_entry(price)