codebase_index 0.1.0

data/lib/codebase_index/railtie.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module CodebaseIndex
+  # Railtie integrates CodebaseIndex into Rails applications.
+  # Loads rake tasks automatically when the gem is bundled.
+  # Conditionally inserts session tracer middleware when enabled.
+  class Railtie < Rails::Railtie
+    rake_tasks do
+      load File.expand_path('../tasks/codebase_index.rake', __dir__)
+    end
+
+    initializer 'codebase_index.session_tracer' do |app|
+      config = CodebaseIndex.configuration
+      if config.session_tracer_enabled
+        require 'codebase_index/session_tracer/middleware'
+
+        app.middleware.use(
+          CodebaseIndex::SessionTracer::Middleware,
+          store: config.session_store,
+          session_id_proc: config.session_id_proc,
+          exclude_paths: config.session_exclude_paths
+        )
+      end
+    end
+  end
+end

data/lib/codebase_index/resilience/circuit_breaker.rb

@@ -0,0 +1,99 @@
+# frozen_string_literal: true
+
+module CodebaseIndex
+  module Resilience
+    # Raised when the circuit breaker is open and calls are being rejected.
+    #
+    # @example Handling a circuit open condition
+    #   begin
+    #     breaker.call { provider.embed(text) }
+    #   rescue CircuitOpenError => e
+    #     use_cached_result(text)
+    #   end
+    class CircuitOpenError < CodebaseIndex::Error; end
+
+    # Circuit breaker pattern for protecting external service calls.
+    #
+    # Tracks failures and transitions between three states:
+    # - **:closed** — normal operation, calls pass through
+    # - **:open** — too many failures, calls are rejected immediately
+    # - **:half_open** — testing recovery, one call is allowed through
+    #
+    # @example Basic usage
+    #   breaker = CircuitBreaker.new(threshold: 5, reset_timeout: 60)
+    #   result = breaker.call { external_service.request }
+    #
+    # @example With retry logic
+    #   breaker = CircuitBreaker.new(threshold: 3, reset_timeout: 30)
+    #   begin
+    #     breaker.call { api.embed(text) }
+    #   rescue CircuitOpenError
+    #     # Service is down, use fallback
+    #   end
+    class CircuitBreaker
+      # @return [Symbol] Current state — :closed, :open, or :half_open
+      attr_reader :state
+
+      # @param threshold [Integer] Number of consecutive failures before opening the circuit
+      # @param reset_timeout [Numeric] Seconds to wait before transitioning from open to half_open
+      def initialize(threshold: 5, reset_timeout: 60)
+        @threshold = threshold
+        @reset_timeout = reset_timeout
+        @state = :closed
+        @failure_count = 0
+        @last_failure_time = nil
+        @mutex = Mutex.new
+      end
+
+      # Execute a block through the circuit breaker.
+      #
+      # @yield The block to execute
+      # @return [Object] The return value of the block
+      # @raise [CircuitOpenError] if the circuit is open and the timeout has not elapsed
+      # @raise [StandardError] re-raises any error from the block
+      def call(&block)
+        # Phase 1: Check state under mutex
+        @mutex.synchronize do
+          case @state
+          when :open
+            unless Time.now - @last_failure_time >= @reset_timeout
+              raise CircuitOpenError, "Circuit breaker is open (#{@failure_count} failures)"
+            end
+
+            @state = :half_open
+          end
+        end
+
+        # Phase 2: Execute outside mutex
+        result = block.call
+
+        # Phase 3: Record success under mutex
+        @mutex.synchronize { reset! }
+
+        result
+      rescue CircuitOpenError
+        raise
+      rescue StandardError => e
+        # Phase 4: Record failure under mutex
+        @mutex.synchronize { record_failure }
+        raise e
+      end
+
+      private
+
+      # Record a failure and potentially open the circuit.
+      def record_failure
+        @failure_count += 1
+        @last_failure_time = Time.now
+        @state = :open if @failure_count >= @threshold
+      end
+
+      # Reset the circuit breaker to closed state with zero failures.
+      def reset!
+        @state = :closed
+        @failure_count = 0
+        @last_failure_time = nil
+      end
+    end
+  end
+end

data/lib/codebase_index/resilience/index_validator.rb

@@ -0,0 +1,185 @@
+# frozen_string_literal: true
+
+require 'json'
+require 'digest'
+
+module CodebaseIndex
+  module Resilience
+    # Validates the integrity of a codebase index output directory.
+    #
+    # Checks that:
+    # - Each type directory has a valid `_index.json`
+    # - All files referenced in the index exist on disk
+    # - Content hashes (source_hash) match the actual source_code
+    # - No stale unit files exist that aren't listed in the index
+    #
+    # @example
+    #   validator = IndexValidator.new(index_dir: "tmp/codebase_index")
+    #   report = validator.validate
+    #   puts report.errors if !report.valid?
+    class IndexValidator
+      # Report produced by {#validate}.
+      #
+      # @!attribute [r] valid?
+      #   @return [Boolean] true if no errors were found
+      # @!attribute [r] warnings
+      #   @return [Array<String>] non-fatal issues (e.g., stale files)
+      # @!attribute [r] errors
+      #   @return [Array<String>] fatal integrity issues
+      ValidationReport = Struct.new(:valid?, :warnings, :errors, keyword_init: true)
+
+      # @param index_dir [String] Path to the codebase index output directory
+      def initialize(index_dir:)
+        @index_dir = index_dir
+      end
+
+      # Validate the index directory and return a report.
+      #
+      # @return [ValidationReport] the validation results
+      def validate
+        warnings = []
+        errors = []
+
+        unless Dir.exist?(@index_dir)
+          errors << "Index directory does not exist: #{@index_dir}"
+          return ValidationReport.new(valid?: false, warnings: warnings, errors: errors)
+        end
+
+        type_dirs = Dir.children(@index_dir).filter_map do |name|
+          full_path = File.join(@index_dir, name)
+          full_path if File.directory?(full_path)
+        end
+
+        type_dirs.each do |type_dir|
+          validate_type_directory(type_dir, warnings, errors)
+        end
+
+        ValidationReport.new(valid?: errors.empty?, warnings: warnings, errors: errors)
+      end
+
+      private
+
+      # Validate a single type directory (e.g., models/, controllers/).
+      #
+      # @param type_dir [String] Absolute path to the type directory
+      # @param warnings [Array<String>] Accumulated warnings
+      # @param errors [Array<String>] Accumulated errors
+      def validate_type_directory(type_dir, warnings, errors)
+        type_name = File.basename(type_dir)
+        index_path = File.join(type_dir, '_index.json')
+
+        unless File.exist?(index_path)
+          errors << "Missing _index.json in #{type_name}/"
+          return
+        end
+
+        index_entries = JSON.parse(File.read(index_path))
+        indexed_identifiers = Set.new
+
+        index_entries.each do |entry|
+          identifier = entry['identifier']
+          indexed_identifiers << identifier
+          validate_index_entry(type_dir, type_name, identifier, errors)
+        end
+
+        check_stale_files(type_dir, type_name, indexed_identifiers, warnings)
+      end
+
+      # Validate that a single index entry has a corresponding unit file with correct hash.
+      #
+      # @param type_dir [String] Path to the type directory
+      # @param type_name [String] Name of the type (for error messages)
+      # @param identifier [String] The unit identifier from the index
+      # @param errors [Array<String>] Accumulated errors
+      def validate_index_entry(type_dir, type_name, identifier, errors)
+        unit_file = find_unit_file(type_dir, identifier)
+
+        unless unit_file
+          errors << "Missing unit file for #{identifier} in #{type_name}/"
+          return
+        end
+
+        validate_content_hash(unit_file, identifier, errors)
+      end
+
+      # Find the JSON file for a given identifier in a type directory.
+      #
+      # @param type_dir [String] Path to the type directory
+      # @param identifier [String] The unit identifier
+      # @return [String, nil] Path to the unit file, or nil if not found
+      def find_unit_file(type_dir, identifier)
+        # Try collision-safe first (current format), then legacy safe_filename, then exact match
+        candidates = [
+          File.join(type_dir, collision_safe_filename(identifier)),
+          File.join(type_dir, safe_filename(identifier)),
+          File.join(type_dir, "#{identifier}.json")
+        ]
+
+        candidates.find { |path| File.exist?(path) }
+      end
+
+      # Validate that the source_hash in a unit file matches the actual source_code.
+      #
+      # @param unit_file [String] Path to the unit JSON file
+      # @param identifier [String] The unit identifier (for error messages)
+      # @param errors [Array<String>] Accumulated errors
+      def validate_content_hash(unit_file, identifier, errors)
+        data = JSON.parse(File.read(unit_file))
+        source_code = data['source_code']
+        stored_hash = data['source_hash']
+
+        return unless source_code && stored_hash
+
+        expected_hash = Digest::SHA256.hexdigest(source_code)
+        return if stored_hash == expected_hash
+
+        errors << "Content hash mismatch for #{identifier}: expected #{expected_hash[0..7]}..., " \
+                  "got #{stored_hash[0..7]}..."
+      end
+
+      # Check for unit files that exist on disk but aren't referenced in the index.
+      #
+      # @param type_dir [String] Path to the type directory
+      # @param type_name [String] Name of the type (for warning messages)
+      # @param indexed_identifiers [Set<String>] Identifiers listed in the index
+      # @param warnings [Array<String>] Accumulated warnings
+      def check_stale_files(type_dir, type_name, indexed_identifiers, warnings)
+        # Build a set of expected filenames from indexed identifiers (both current and legacy formats)
+        expected_filenames = Set.new
+        indexed_identifiers.each do |id|
+          expected_filenames << collision_safe_filename(id)
+          expected_filenames << safe_filename(id)
+          expected_filenames << "#{id}.json"
+        end
+
+        Dir[File.join(type_dir, '*.json')].each do |file|
+          basename = File.basename(file)
+          next if basename == '_index.json'
+          next if expected_filenames.include?(basename)
+
+          warnings << "Stale file not in index: #{type_name}/#{basename}"
+        end
+      end
+
+      # Convert an identifier to a safe filename (legacy format, mirrors Extractor#safe_filename).
+      #
+      # @param identifier [String] The unit identifier (e.g., "Admin::UsersController")
+      # @return [String] A filesystem-safe filename (e.g., "Admin__UsersController.json")
+      def safe_filename(identifier)
+        "#{identifier.gsub('::', '__').gsub(/[^a-zA-Z0-9_-]/, '_')}.json"
+      end
+
+      # Convert an identifier to a collision-safe filename (current format).
+      # Mirrors {Extractor#collision_safe_filename} — appends a short SHA256 digest
+      # to disambiguate identifiers that normalize to the same safe_filename.
+      #
+      # @param identifier [String] The unit identifier
+      # @return [String] Collision-safe filename (e.g., "Admin__UsersController_a1b2c3d4.json")
+      def collision_safe_filename(identifier)
+        base = identifier.gsub('::', '__').gsub(/[^a-zA-Z0-9_-]/, '_')
+        digest = Digest::SHA256.hexdigest(identifier)[0, 8]
+        "#{base}_#{digest}.json"
+      end
+    end
+  end
+end

data/lib/codebase_index/resilience/retryable_provider.rb

@@ -0,0 +1,108 @@
+# frozen_string_literal: true
+
+require_relative '../embedding/provider'
+require_relative 'circuit_breaker'
+
+module CodebaseIndex
+  module Resilience
+    # Wraps an embedding provider with retry logic and optional circuit breaker.
+    #
+    # Transparently retries transient failures with exponential backoff.
+    # When a circuit breaker is provided, all calls are routed through it,
+    # and {CircuitOpenError} is never retried.
+    #
+    # @example Without circuit breaker
+    #   retryable = RetryableProvider.new(provider: ollama_provider, max_retries: 3)
+    #   vector = retryable.embed("some text")
+    #
+    # @example With circuit breaker
+    #   breaker = CircuitBreaker.new(threshold: 5, reset_timeout: 60)
+    #   retryable = RetryableProvider.new(
+    #     provider: ollama_provider,
+    #     max_retries: 3,
+    #     circuit_breaker: breaker
+    #   )
+    #   vector = retryable.embed("some text")
+    class RetryableProvider
+      include CodebaseIndex::Embedding::Provider::Interface
+
+      # @param provider [#embed, #embed_batch, #dimensions, #model_name] The underlying embedding provider
+      # @param max_retries [Integer] Maximum number of retry attempts
+      # @param circuit_breaker [CircuitBreaker, nil] Optional circuit breaker instance
+      def initialize(provider:, max_retries: 3, circuit_breaker: nil)
+        @provider = provider
+        @max_retries = max_retries
+        @circuit_breaker = circuit_breaker
+      end
+
+      # Embed a single text string with retry logic.
+      #
+      # @param text [String] the text to embed
+      # @return [Array<Float>] the embedding vector
+      # @raise [CircuitOpenError] if the circuit breaker is open
+      # @raise [StandardError] if all retries are exhausted
+      def embed(text)
+        with_retries { call_provider { @provider.embed(text) } }
+      end
+
+      # Embed multiple texts with retry logic.
+      #
+      # @param texts [Array<String>] the texts to embed
+      # @return [Array<Array<Float>>] array of embedding vectors
+      # @raise [CircuitOpenError] if the circuit breaker is open
+      # @raise [StandardError] if all retries are exhausted
+      def embed_batch(texts)
+        with_retries { call_provider { @provider.embed_batch(texts) } }
+      end
+
+      # Return the dimensionality of the embedding vectors.
+      #
+      # @return [Integer] number of dimensions
+      def dimensions
+        @provider.dimensions
+      end
+
+      # Return the name of the embedding model.
+      #
+      # @return [String] model name
+      def model_name
+        @provider.model_name
+      end
+
+      private
+
+      # Execute a block with retry logic and exponential backoff.
+      #
+      # @yield The block to execute
+      # @return [Object] The return value of the block
+      # @raise [CircuitOpenError] immediately without retrying
+      # @raise [StandardError] the last error if all retries are exhausted
+      def with_retries
+        attempt = 0
+        begin
+          attempt += 1
+          yield
+        rescue CircuitOpenError
+          raise
+        rescue StandardError => e
+          raise e if attempt > @max_retries
+
+          sleep((2**attempt) * 0.1)
+          retry
+        end
+      end
+
+      # Route a call through the circuit breaker if one is configured.
+      #
+      # @yield The block to execute
+      # @return [Object] The return value of the block
+      def call_provider(&block)
+        if @circuit_breaker
+          @circuit_breaker.call(&block)
+        else
+          block.call
+        end
+      end
+    end
+  end
+end

data/lib/codebase_index/retrieval/context_assembler.rb

@@ -0,0 +1,249 @@
+# frozen_string_literal: true
+
+module CodebaseIndex
+  module Retrieval
+    # Transforms ranked search candidates into a token-budgeted context string
+    # for LLM consumption.
+    #
+    # Allocates a fixed token budget across four sections:
+    # - Structural (10%): Always-included codebase overview
+    # - Primary (50%): Direct query results
+    # - Supporting (25%): Dependencies and related context
+    # - Framework (15%): Rails/gem source when query has framework context
+    #
+    # When framework context is not needed, primary and supporting sections
+    # receive the framework allocation proportionally.
+    #
+    # @example
+    #   assembler = ContextAssembler.new(metadata_store: store)
+    #   result = assembler.assemble(candidates: ranked, classification: cls)
+    #   result.context     # => "## User (model)\n..."
+    #   result.tokens_used # => 4200
+    #   result.sections    # => [:structural, :primary, :supporting]
+    #
+    class ContextAssembler
+      DEFAULT_BUDGET = 8000 # tokens
+
+      BUDGET_ALLOCATION = {
+        structural: 0.10,
+        primary: 0.50,
+        supporting: 0.25,
+        framework: 0.15
+      }.freeze
+
+      # Minimum token count for a section to be worth including.
+      MIN_USEFUL_TOKENS = 200
+
+      # @param metadata_store [#find] Store that resolves identifiers to unit data
+      # @param budget [Integer] Total token budget
+      def initialize(metadata_store:, budget: DEFAULT_BUDGET)
+        @metadata_store = metadata_store
+        @budget = budget
+      end
+
+      # Assemble context from ranked candidates within token budget.
+      #
+      # @param candidates [Array<Candidate>] Ranked search candidates
+      # @param classification [QueryClassifier::Classification] Query classification
+      # @param structural_context [String, nil] Optional codebase overview text
+      # @param budget [Integer, nil] Override token budget; falls back to @budget
+      # @return [AssembledContext] Token-budgeted context with source attribution
+      def assemble(candidates:, classification:, structural_context: nil, budget: nil)
+        effective_budget = budget || @budget
+        sections = []
+        sources = []
+        tokens_used = 0
+
+        # 1. Structural context (always first if provided)
+        tokens_used = add_structural_section(sections, structural_context, tokens_used, effective_budget)
+
+        # 2. Compute per-section budgets from remaining tokens
+        budgets = compute_section_budgets(effective_budget - tokens_used, classification)
+
+        # 3. Primary, supporting, and framework sections
+        add_candidate_section(sections, sources, :primary,
+                              candidates.reject { |c| c.source == :graph_expansion }, budgets[:primary])
+        add_candidate_section(sections, sources, :supporting,
+                              candidates.select { |c| c.source == :graph_expansion }, budgets[:supporting])
+        if budgets[:framework].positive?
+          add_candidate_section(sections, sources, :framework,
+                                candidates.select { |c| framework_candidate?(c) }, budgets[:framework])
+        end
+
+        build_result(sections, sources, effective_budget)
+      end
+
+      private
+
+      # Add structural context section if provided.
+      #
+      # @return [Integer] Updated tokens_used count
+      def add_structural_section(sections, structural_context, tokens_used, effective_budget)
+        return tokens_used unless structural_context
+
+        budget = (effective_budget * BUDGET_ALLOCATION[:structural]).to_i
+        text = truncate_to_budget(structural_context, budget)
+        sections << { section: :structural, content: text }
+        tokens_used + estimate_tokens(text)
+      end
+
+      # Add a candidate-based section if candidates produce content.
+      #
+      # @return [void]
+      def add_candidate_section(sections, sources, section_name, candidates, budget)
+        return if candidates.empty?
+
+        content, section_sources = assemble_section(candidates, budget)
+        return if content.empty?
+
+        sections << { section: section_name, content: content }
+        sources.concat(section_sources)
+      end
+
+      # Compute token budgets for primary/supporting/framework sections.
+      #
+      # @param remaining [Integer] Tokens available after structural
+      # @param classification [QueryClassifier::Classification]
+      # @return [Hash<Symbol, Integer>]
+      def compute_section_budgets(remaining, classification)
+        if classification.framework_context
+          {
+            primary: (remaining * 0.55).to_i,
+            supporting: (remaining * 0.25).to_i,
+            framework: (remaining * 0.20).to_i
+          }
+        else
+          {
+            primary: (remaining * 0.65).to_i,
+            supporting: (remaining * 0.35).to_i,
+            framework: 0
+          }
+        end
+      end
+
+      # Assemble content for a single section within a token budget.
+      #
+      # @param candidates [Array<Candidate>] Candidates for this section
+      # @param budget [Integer] Token budget for this section
+      # @return [Array(String, Array<Hash>)] Content string and source attributions
+      def assemble_section(candidates, budget)
+        content_parts = []
+        sources = []
+        tokens_used = 0
+
+        candidates.sort_by { |c| -c.score }.each do |candidate|
+          tokens_used = append_candidate(content_parts, sources, candidate, budget, tokens_used)
+          break if tokens_used.nil?
+        end
+
+        [content_parts.join("\n\n"), sources]
+      end
+
+      # Append a single candidate to the section. Returns updated tokens_used, or nil to stop.
+      def append_candidate(parts, sources, candidate, budget, tokens_used)
+        unit = @metadata_store.find(candidate.identifier)
+        return tokens_used unless unit
+
+        text = format_unit(unit, candidate)
+        tokens = estimate_tokens(text)
+        remaining = budget - tokens_used
+
+        if tokens <= remaining
+          parts << text
+          sources << build_source_attribution(candidate, unit)
+          tokens_used + tokens
+        elsif remaining > MIN_USEFUL_TOKENS
+          parts << truncate_to_budget(text, remaining)
+          sources << build_source_attribution(candidate, unit, truncated: true)
+          nil
+        end
+      end
+
+      # Format a unit for inclusion in context.
+      #
+      # @param unit [Hash] Unit data from metadata store
+      # @param candidate [Candidate] The search candidate
+      # @return [String]
+      def format_unit(unit, _candidate)
+        identifier = unit[:identifier] || unit['identifier']
+        type = unit[:type] || unit['type']
+        file_path = unit[:file_path] || unit['file_path']
+        source = unit[:source_code] || unit['source_code'] || ''
+
+        <<~UNIT.strip
+          ## #{identifier} (#{type})
+          File: #{file_path}
+
+          #{source}
+        UNIT
+      end
+
+      # Build source attribution hash for a candidate.
+      #
+      # @return [Hash]
+      def build_source_attribution(candidate, unit, truncated: false)
+        attribution = {
+          identifier: candidate.identifier,
+          type: unit[:type] || unit['type'],
+          score: candidate.score,
+          file_path: unit[:file_path] || unit['file_path']
+        }
+        attribution[:truncated] = true if truncated
+        attribution
+      end
+
+      # Check if a candidate is framework source.
+      #
+      # @param candidate [Candidate]
+      # @return [Boolean]
+      def framework_candidate?(candidate)
+        metadata = candidate.metadata
+        return false unless metadata
+
+        type = metadata[:type] || metadata['type']
+        %w[rails_source gem_source].include?(type.to_s)
+      end
+
+      # Truncate text to fit within a token budget.
+      #
+      # @param text [String]
+      # @param token_budget [Integer]
+      # @return [String]
+      def truncate_to_budget(text, token_budget)
+        return text if estimate_tokens(text) <= token_budget
+
+        # Estimate target character count with 10% safety margin
+        target_chars = (token_budget * 4.0 * 0.9).to_i
+        "#{text[0...target_chars]}\n... [truncated]"
+      end
+
+      # Estimate token count using the project convention.
+      #
+      # @param text [String]
+      # @return [Integer]
+      def estimate_tokens(text)
+        (text.length / 4.0).ceil
+      end
+
+      # Build the final AssembledContext result.
+      #
+      # @param sections [Array<Hash>] Assembled sections
+      # @param sources [Array<Hash>] Source attributions
+      # @param effective_budget [Integer] The budget actually used for assembly
+      # @return [AssembledContext]
+      def build_result(sections, sources, effective_budget)
+        context = sections.map { |s| s[:content] }.join("\n\n---\n\n")
+        AssembledContext.new(
+          context: context,
+          tokens_used: estimate_tokens(context),
+          budget: effective_budget,
+          sources: sources.uniq,
+          sections: sections.map { |s| s[:section] }
+        )
+      end
+    end
+
+    # Result of context assembly.
+    AssembledContext = Struct.new(:context, :tokens_used, :budget, :sources, :sections, keyword_init: true)
+  end
+end