rubyn-code 0.2.2 → 0.3.0

data/lib/rubyn_code/mcp/tool_bridge.rb

@@ -11,6 +11,11 @@ module RubynCode
     # - Delegates #execute to the MCP client's #call_tool
     # - Registers itself with Tools::Registry
     module ToolBridge
+      JSON_TYPE_MAP = {
+        'string' => :string, 'integer' => :integer, 'number' => :number,
+        'boolean' => :boolean, 'array' => :array, 'object' => :object
+      }.freeze
+
       class << self
         # Discovers tools from an MCP client and creates corresponding
         # RubynCode tool classes.
@@ -35,12 +40,21 @@ module RubynCode
         # @return [Class] the newly created and registered tool class
         def build_tool_class(mcp_client, tool_def)
           remote_name = tool_def['name']
-          tool_name = "mcp_#{sanitize_name(remote_name)}"
-          description = tool_def['description'] || "MCP tool: #{remote_name}"
-          input_schema = tool_def['inputSchema'] || {}
-          parameters = build_parameters_from_schema(input_schema)
+          attrs = {
+            tool_name: "mcp_#{sanitize_name(remote_name)}",
+            description: tool_def['description'] || "MCP tool: #{remote_name}",
+            parameters: build_parameters_from_schema(tool_def['inputSchema'] || {})
+          }
+          klass = create_tool_class(attrs[:tool_name], attrs[:description], attrs[:parameters], mcp_client,
+                                    remote_name)
+          Tools::Registry.register(klass)
+          klass
+        end
 
-          klass = Class.new(Tools::Base) do
+        def create_tool_class(tool_name, description, parameters, mcp_client, remote_name) # rubocop:disable Metrics/AbcSize, Metrics/MethodLength -- dynamic class creation requires setting many constants
+          bridge = self
+
+          Class.new(Tools::Base) do
             const_set(:TOOL_NAME, tool_name)
             const_set(:DESCRIPTION, description)
             const_set(:PARAMETERS, parameters)
@@ -60,62 +74,25 @@ module RubynCode
             define_method(:format_result) do |result|
               case result
               when Hash
-                if result.key?('content')
-                  extract_content(result['content'])
-                else
-                  JSON.generate(result)
-                end
-              when String
-                result
-              else
-                result.to_s
+                result.key?('content') ? extract_content(result['content']) : JSON.generate(result)
+              when String then result
+              else result.to_s
               end
             end
 
             define_method(:extract_content) do |content|
-              Array(content).map do |block|
-                case block['type']
-                when 'text'
-                  block['text']
-                when 'image'
-                  "[image: #{block['mimeType']}]"
-                when 'resource'
-                  block.dig('resource', 'text') || "[resource: #{block.dig('resource', 'uri')}]"
-                else
-                  block.to_s
-                end
-              end.join("\n")
-            end
-          end
-
-          # Build parameter definitions from JSON Schema
-          klass.define_singleton_method(:build_parameters) do |schema|
-            properties = schema['properties'] || {}
-            required = schema['required'] || []
-
-            properties.each_with_object({}) do |(name, prop), params|
-              params[name.to_sym] = {
-                type: map_json_type(prop['type']),
-                description: prop['description'] || '',
-                required: required.include?(name)
-              }
+              Array(content).map { |block| bridge.send(:format_mcp_block, block) }.join("\n")
             end
           end
+        end
 
-          klass.define_singleton_method(:map_json_type) do |json_type|
-            case json_type
-            when 'string'  then :string
-            when 'integer' then :integer
-            when 'number'  then :number
-            when 'boolean' then :boolean
-            when 'array'   then :array
-            when 'object'  then :object
-            else :string
-            end
+        def format_mcp_block(block)
+          case block['type']
+          when 'text'     then block['text']
+          when 'image'    then "[image: #{block['mimeType']}]"
+          when 'resource' then block.dig('resource', 'text') || "[resource: #{block.dig('resource', 'uri')}]"
+          else block.to_s
           end
-
-          Tools::Registry.register(klass)
-          klass
         end
 
         # Builds parameter definitions from a JSON Schema.
@@ -140,15 +117,7 @@ module RubynCode
         # @param json_type [String]
         # @return [Symbol]
         def map_json_type(json_type)
-          case json_type
-          when 'string'  then :string
-          when 'integer' then :integer
-          when 'number'  then :number
-          when 'boolean' then :boolean
-          when 'array'   then :array
-          when 'object'  then :object
-          else :string
-          end
+          JSON_TYPE_MAP.fetch(json_type, :string)
         end
 
         # Sanitizes a tool name for use as a Ruby-friendly identifier.

data/lib/rubyn_code/memory/session_persistence.rb

@@ -16,19 +16,21 @@ module RubynCode
 
       # Persists a complete session snapshot.
       #
-      # @param session_id [String] unique session identifier
-      # @param project_path [String] project this session belongs to
-      # @param messages [Array<Hash>] the conversation messages
-      # @param title [String, nil] human-readable session title
-      # @param model [String, nil] LLM model used
-      # @param metadata [Hash] arbitrary metadata
+      # @param attrs [Hash] session attributes:
+      #   :session_id, :project_path, :messages (required);
+      #   :title, :model, :metadata (optional)
       # @return [void]
-      def save_session(session_id:, project_path:, messages:, title: nil, model: nil, metadata: {})
+      def save_session(session_id:, project_path:, messages:, **opts)
         now = Time.now.utc.strftime('%Y-%m-%d %H:%M:%S')
         messages_json = JSON.generate(messages)
-        meta_json = JSON.generate(metadata)
+        meta_json = JSON.generate(opts.fetch(:metadata, {}))
+        title = opts[:title]
+        model = opts[:model]
 
-        @db.execute(<<~SQL, [session_id, project_path, title, model, messages_json, 'active', meta_json, now, now, messages_json, title, model, meta_json, now])
+        insert_params = [session_id, project_path, title, model, messages_json, 'active', meta_json, now, now]
+        update_params = [messages_json, title, model, meta_json, now]
+
+        @db.execute(<<~SQL, insert_params + update_params)
           INSERT INTO sessions (id, project_path, title, model, messages, status, metadata, created_at, updated_at)
           VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
           ON CONFLICT(id) DO UPDATE SET
@@ -71,20 +73,7 @@ module RubynCode
       # @param limit [Integer] maximum results (default 20)
       # @return [Array<Hash>] session summaries (without full messages)
       def list_sessions(project_path: nil, status: nil, limit: 20)
-        conditions = []
-        params = []
-
-        if project_path
-          conditions << 'project_path = ?'
-          params << project_path
-        end
-
-        if status
-          conditions << 'status = ?'
-          params << status
-        end
-
-        where_clause = conditions.empty? ? '' : "WHERE #{conditions.join(' AND ')}"
+        where_clause, params = build_list_filters(project_path, status)
         params << limit
 
         rows = @db.query(<<~SQL, params).to_a
@@ -95,18 +84,7 @@ module RubynCode
           LIMIT ?
         SQL
 
-        rows.map do |row|
-          {
-            id: row['id'],
-            project_path: row['project_path'],
-            title: row['title'],
-            model: row['model'],
-            status: row['status'],
-            metadata: parse_json_hash(row['metadata']),
-            created_at: row['created_at'],
-            updated_at: row['updated_at']
-          }
-        end
+        rows.map { |row| row_to_session_summary(row) }
       end
 
       # Updates session attributes.
@@ -117,29 +95,7 @@ module RubynCode
       def update_session(session_id, **attrs)
         return if attrs.empty?
 
-        sets = []
-        params = []
-
-        attrs.each do |key, value|
-          case key
-          when :title
-            sets << 'title = ?'
-            params << value
-          when :status
-            sets << 'status = ?'
-            params << value
-          when :model
-            sets << 'model = ?'
-            params << value
-          when :metadata
-            sets << 'metadata = ?'
-            params << JSON.generate(value)
-          when :messages
-            sets << 'messages = ?'
-            params << JSON.generate(value)
-          end
-        end
-
+        sets, params = build_update_clauses(attrs)
         return if sets.empty?
 
         sets << 'updated_at = ?'
@@ -157,8 +113,53 @@ module RubynCode
         @db.execute('DELETE FROM sessions WHERE id = ?', [session_id])
       end
 
+      JSON_ATTRS = %i[metadata messages].freeze
+      SIMPLE_ATTRS = %i[title status model].freeze
+
       private
 
+      def build_list_filters(project_path, status)
+        conditions = []
+        params = []
+        if project_path
+          conditions << 'project_path = ?'
+          params << project_path
+        end
+        if status
+          conditions << 'status = ?'
+          params << status
+        end
+        where_clause = conditions.empty? ? '' : "WHERE #{conditions.join(' AND ')}"
+        [where_clause, params]
+      end
+
+      def row_to_session_summary(row)
+        {
+          id: row['id'],
+          project_path: row['project_path'],
+          title: row['title'],
+          model: row['model'],
+          status: row['status'],
+          metadata: parse_json_hash(row['metadata']),
+          created_at: row['created_at'],
+          updated_at: row['updated_at']
+        }
+      end
+
+      def build_update_clauses(attrs)
+        sets = []
+        params = []
+
+        attrs.each do |key, value|
+          next unless SIMPLE_ATTRS.include?(key) || JSON_ATTRS.include?(key)
+
+          sets << "#{key} = ?"
+          params << (JSON_ATTRS.include?(key) ? JSON.generate(value) : value)
+        end
+
+        [sets, params]
+      end
+
       def ensure_table
         @db.execute(<<~SQL)
           CREATE TABLE IF NOT EXISTS sessions (

data/lib/rubyn_code/memory/store.rb

@@ -57,34 +57,7 @@ module RubynCode
       def update(id, **attrs)
         return if attrs.empty?
 
-        sets = []
-        params = []
-
-        attrs.each do |key, value|
-          case key
-          when :content
-            sets << 'content = ?'
-            params << value
-          when :tier
-            validate_tier!(value)
-            sets << 'tier = ?'
-            params << value
-          when :category
-            validate_category!(value) if value
-            sets << 'category = ?'
-            params << value
-          when :metadata
-            sets << 'metadata = ?'
-            params << JSON.generate(value)
-          when :expires_at
-            sets << 'expires_at = ?'
-            params << value
-          when :relevance_score
-            sets << 'relevance_score = ?'
-            params << value.to_f
-          end
-        end
-
+        sets, params = build_memory_update(attrs)
         return if sets.empty?
 
         params << id
@@ -92,8 +65,6 @@ module RubynCode
           "UPDATE memories SET #{sets.join(', ')} WHERE id = ? AND project_path = '#{@project_path}'",
           params
         )
-
-        # Content changes are picked up by LIKE-based search — no FTS sync needed
       end
 
       # Deletes a memory and its FTS index entry.
@@ -142,41 +113,57 @@ module RubynCode
         SQL
       end
 
+      MEMORY_ATTR_MAP = {
+        content: ->(v) { v },
+        tier: ->(v) { v },
+        category: ->(v) { v },
+        metadata: ->(v) { JSON.generate(v) },
+        expires_at: ->(v) { v },
+        relevance_score: lambda(&:to_f)
+      }.freeze
+
       private
 
+      def build_memory_update(attrs)
+        sets = []
+        params = []
+
+        attrs.each do |key, value|
+          next unless MEMORY_ATTR_MAP.key?(key)
+
+          validate_tier!(value) if key == :tier
+          validate_category!(value) if key == :category && value
+
+          sets << "#{key} = ?"
+          params << MEMORY_ATTR_MAP[key].call(value)
+        end
+
+        [sets, params]
+      end
+
       def ensure_tables
+        create_memories_table
+        create_memories_indexes
+      end
+
+      def create_memories_table
         @db.execute(<<~SQL)
           CREATE TABLE IF NOT EXISTS memories (
-            id              TEXT PRIMARY KEY,
-            project_path    TEXT NOT NULL,
-            tier            TEXT NOT NULL DEFAULT 'medium',
-            category        TEXT,
-            content         TEXT NOT NULL,
-            relevance_score REAL NOT NULL DEFAULT 1.0,
-            access_count    INTEGER NOT NULL DEFAULT 0,
-            last_accessed_at TEXT,
-            expires_at      TEXT,
-            metadata        TEXT DEFAULT '{}',
-            created_at      TEXT NOT NULL
+            id TEXT PRIMARY KEY, project_path TEXT NOT NULL,
+            tier TEXT NOT NULL DEFAULT 'medium', category TEXT,
+            content TEXT NOT NULL, relevance_score REAL NOT NULL DEFAULT 1.0,
+            access_count INTEGER NOT NULL DEFAULT 0, last_accessed_at TEXT,
+            expires_at TEXT, metadata TEXT DEFAULT '{}', created_at TEXT NOT NULL
           )
         SQL
+      end
 
+      def create_memories_indexes
+        @db.execute('CREATE INDEX IF NOT EXISTS idx_memories_project_tier ON memories (project_path, tier)')
+        @db.execute('CREATE INDEX IF NOT EXISTS idx_memories_project_category ON memories (project_path, category)')
         @db.execute(<<~SQL)
-          CREATE INDEX IF NOT EXISTS idx_memories_project_tier
-          ON memories (project_path, tier)
-        SQL
-
-        @db.execute(<<~SQL)
-          CREATE INDEX IF NOT EXISTS idx_memories_project_category
-          ON memories (project_path, category)
-        SQL
-
-        @db.execute(<<~SQL)
-          CREATE INDEX IF NOT EXISTS idx_memories_expires_at
-          ON memories (expires_at) WHERE expires_at IS NOT NULL
+          CREATE INDEX IF NOT EXISTS idx_memories_expires_at ON memories (expires_at) WHERE expires_at IS NOT NULL
         SQL
-
-        # Search uses LIKE queries — no FTS table needed
       end
 
       # @param tier [String]

data/lib/rubyn_code/observability/budget_enforcer.rb

@@ -37,39 +37,17 @@ module RubynCode
       # @param cache_write_tokens [Integer] cache-write token count
       # @param request_type [String] the type of request (e.g., "chat", "compact")
       # @return [CostRecord] the persisted cost record
-      def record!(model:, input_tokens:, output_tokens:, cache_read_tokens: 0, cache_write_tokens: 0,
-                  request_type: 'chat')
-        cost = CostCalculator.calculate(
-          model: model,
-          input_tokens: input_tokens,
-          output_tokens: output_tokens,
-          cache_read_tokens: cache_read_tokens,
-          cache_write_tokens: cache_write_tokens
-        )
-
-        id = SecureRandom.uuid
-        now = Time.now.utc.iso8601
+      def record!(model:, input_tokens:, output_tokens:, **opts)
+        cache_read  = opts.fetch(:cache_read_tokens, 0)
+        cache_write = opts.fetch(:cache_write_tokens, 0)
+        req_type    = opts.fetch(:request_type, 'chat')
 
-        @db.execute(
-          "INSERT INTO #{TABLE_NAME} (id, session_id, model, input_tokens, output_tokens, " \
-          'cache_read_tokens, cache_write_tokens, cost_usd, request_type, created_at) ' \
-          'VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)',
-          [id, @session_id, model, input_tokens, output_tokens,
-           cache_read_tokens, cache_write_tokens, cost, request_type, now]
+        cost = CostCalculator.calculate(
+          model: model, input_tokens: input_tokens, output_tokens: output_tokens,
+          cache_read_tokens: cache_read, cache_write_tokens: cache_write
         )
 
-        CostRecord.new(
-          id: id,
-          session_id: @session_id,
-          model: model,
-          input_tokens: input_tokens,
-          output_tokens: output_tokens,
-          cache_read_tokens: cache_read_tokens,
-          cache_write_tokens: cache_write_tokens,
-          cost_usd: cost,
-          request_type: request_type,
-          created_at: now
-        )
+        persist_cost_record(model, input_tokens, output_tokens, cache_read, cache_write, cost, req_type)
       end
 
       # Raises BudgetExceededError if either the session or daily budget is exceeded.
@@ -80,14 +58,16 @@ module RubynCode
         sc = session_cost
         if sc >= @session_limit
           raise BudgetExceededError,
-                "Session budget exceeded: $#{'%.4f' % sc} >= $#{format('%.2f', @session_limit)} limit"
+                format('Session budget exceeded: $%<cost>.4f >= $%<limit>.2f limit',
+                       cost: sc, limit: @session_limit)
         end
 
         dc = daily_cost
         return unless dc >= @daily_limit
 
         raise BudgetExceededError,
-              "Daily budget exceeded: $#{'%.4f' % dc} >= $#{format('%.2f', @daily_limit)} limit"
+              format('Daily budget exceeded: $%<cost>.4f >= $%<limit>.2f limit',
+                     cost: dc, limit: @daily_limit)
       end
 
       # Returns the total cost accumulated in the current session.
@@ -124,6 +104,40 @@ module RubynCode
 
       private
 
+      CostRecordAttrs = Data.define(:model, :input_tokens, :output_tokens, :cache_read, :cache_write, :cost,
+                                    :req_type)
+      private_constant :CostRecordAttrs
+
+      def persist_cost_record(model, input_tokens, output_tokens, cache_read, cache_write, cost, req_type) # rubocop:disable Metrics/ParameterLists -- maps directly to DB columns
+        attrs = CostRecordAttrs.new(model:, input_tokens:, output_tokens:, cache_read:, cache_write:, cost:,
+                                    req_type:)
+        insert_cost_record(attrs)
+      end
+
+      def insert_cost_record(attrs)
+        record_id = SecureRandom.uuid
+        now = Time.now.utc.iso8601
+
+        @db.execute(
+          "INSERT INTO #{TABLE_NAME} (id, session_id, model, input_tokens, output_tokens, " \
+          'cache_read_tokens, cache_write_tokens, cost_usd, request_type, created_at) ' \
+          'VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)',
+          [record_id, @session_id, attrs.model, attrs.input_tokens, attrs.output_tokens,
+           attrs.cache_read, attrs.cache_write, attrs.cost, attrs.req_type, now]
+        )
+
+        build_cost_record(record_id, attrs, now)
+      end
+
+      def build_cost_record(record_id, attrs, now)
+        CostRecord.new(
+          id: record_id, session_id: @session_id, model: attrs.model,
+          input_tokens: attrs.input_tokens, output_tokens: attrs.output_tokens,
+          cache_read_tokens: attrs.cache_read, cache_write_tokens: attrs.cache_write,
+          cost_usd: attrs.cost, request_type: attrs.req_type, created_at: now
+        )
+      end
+
       def ensure_table_exists
         @db.execute(<<~SQL)
           CREATE TABLE IF NOT EXISTS #{TABLE_NAME} (

data/lib/rubyn_code/observability/cost_calculator.rb

@@ -9,9 +9,19 @@ module RubynCode
     module CostCalculator
       # Per-million-token rates: { model_prefix => [input_rate, output_rate] }
       PRICING = {
+        # Anthropic — Claude 4.6
         'claude-haiku-4-5' => [1.00, 5.00],
-        'claude-sonnet-4-20250514' => [3.00, 15.00],
-        'claude-opus-4-20250514' => [15.00, 75.00]
+        'claude-sonnet-4-6' => [3.00, 15.00],
+        'claude-opus-4-6' => [15.00, 75.00],
+        # OpenAI — GPT-5.4
+        'gpt-5.4' => [2.50, 10.00],
+        'gpt-5.4-mini' => [0.15, 0.60],
+        'gpt-5.4-nano' => [0.10, 0.40],
+        # OpenAI — legacy
+        'gpt-4o' => [2.50, 10.00],
+        'gpt-4o-mini' => [0.15, 0.60],
+        'o3' => [2.00, 8.00],
+        'o4-mini' => [1.10, 4.40]
       }.freeze
 
       CACHE_READ_DISCOUNT  = 0.1
@@ -29,12 +39,10 @@ module RubynCode
         def calculate(model:, input_tokens:, output_tokens:, cache_read_tokens: 0, cache_write_tokens: 0)
           input_rate, output_rate = rates_for(model)
 
-          input_cost       = (input_tokens.to_f / 1_000_000) * input_rate
-          output_cost      = (output_tokens.to_f / 1_000_000) * output_rate
-          cache_read_cost  = (cache_read_tokens.to_f / 1_000_000) * input_rate * CACHE_READ_DISCOUNT
-          cache_write_cost = (cache_write_tokens.to_f / 1_000_000) * input_rate * CACHE_WRITE_PREMIUM
-
-          input_cost + output_cost + cache_read_cost + cache_write_cost
+          token_cost(input_tokens, input_rate) +
+            token_cost(output_tokens, output_rate) +
+            token_cost(cache_read_tokens, input_rate * CACHE_READ_DISCOUNT) +
+            token_cost(cache_write_tokens, input_rate * CACHE_WRITE_PREMIUM)
         end
 
         private
@@ -44,7 +52,15 @@ module RubynCode
         #
         # @param model [String]
         # @return [Array(Float, Float)] [input_rate, output_rate]
+        def token_cost(tokens, rate)
+          (tokens.to_f / 1_000_000) * rate
+        end
+
         def rates_for(model)
+          # User-configured pricing takes priority
+          custom = config_pricing(model)
+          return custom if custom
+
           return PRICING[model] if PRICING.key?(model)
 
           # Try prefix match (e.g., "claude-sonnet-4-20250514-v2" matches "claude-sonnet-4-20250514")
@@ -55,6 +71,14 @@ module RubynCode
           # Conservative fallback: use the most expensive known model
           PRICING.max_by { |_, rates| rates.first }.last
         end
+
+        def config_pricing(model)
+          settings = Config::Settings.new
+          custom = settings.custom_pricing
+          custom[model]
+        rescue StandardError
+          nil
+        end
       end
     end
   end