rubyn-code 0.2.2 → 0.4.0

data/lib/rubyn_code/auth/token_store.rb

@@ -7,7 +7,7 @@ require 'time'
 
 module RubynCode
   module Auth
-    module TokenStore
+    module TokenStore # rubocop:disable Metrics/ModuleLength -- single-responsibility credential store
       EXPIRY_BUFFER_SECONDS = 300 # 5 minutes
       KEYCHAIN_SERVICE = 'Claude Code-credentials'
 
@@ -20,82 +20,98 @@ module RubynCode
           load_from_keychain || load_from_file || load_from_env
         end
 
-        def save(access_token:, refresh_token:, expires_at:)
+        # Load API key for a given provider. Anthropic uses the full fallback chain.
+        # Other providers: stored key → env var.
+        def load_for_provider(provider)
+          return load if provider == 'anthropic'
+
+          stored = load_provider_key(provider)
+          return { access_token: stored, type: :api_key, source: :stored } if stored
+
+          env_key = resolve_env_key(provider)
+          api_key = ENV.fetch(env_key, nil)
+          api_key&.empty? == false ? { access_token: api_key, type: :api_key, source: :env } : nil
+        end
+
+        # Store an API key for a provider in tokens.yml (encrypted at rest).
+        def save_provider_key(provider, key)
           ensure_directory!
+          data = load_tokens_file || {}
+          data['provider_keys'] ||= {}
+          data['provider_keys'][provider.to_s] = KeyEncryption.encrypt(key)
+          write_tokens_file(data)
+        end
 
-          data = {
-            'access_token' => access_token,
-            'refresh_token' => refresh_token,
-            'expires_at' => expires_at.is_a?(Time) ? expires_at.iso8601 : expires_at.to_s
-          }
+        # Retrieve a stored API key for a provider (decrypted transparently).
+        def load_provider_key(provider)
+          data = load_tokens_file
+          value = data&.dig('provider_keys', provider.to_s)
+          return nil unless value
 
-          File.write(tokens_path, YAML.dump(data))
-          File.chmod(0o600, tokens_path)
+          migrate_plaintext_key!(data, provider, value) unless KeyEncryption.encrypted?(value)
+          KeyEncryption.decrypt(value)
+        end
+
+        def save(access_token:, refresh_token:, expires_at:)
+          ensure_directory!
+          data = load_tokens_file || {}
+          data['access_token'] = access_token
+          data['refresh_token'] = refresh_token
+          data['expires_at'] = expires_at.is_a?(Time) ? expires_at.iso8601 : expires_at.to_s
+          write_tokens_file(data)
           data
         end
 
-        def clear!
+        def clear! # rubocop:disable Naming/PredicateMethod -- destructive action, not a predicate
           FileUtils.rm_f(tokens_path)
           true
         end
 
         def valid?
           tokens = self.load
-          return false unless tokens
-          return false unless tokens[:access_token]
-
-          # API keys don't expire
+          return false unless tokens&.fetch(:access_token, nil)
           return true if tokens[:type] == :api_key
-
-          # OAuth tokens need expiry check
           return true unless tokens[:expires_at]
 
           tokens[:expires_at] > Time.now + EXPIRY_BUFFER_SECONDS
         end
 
-        def exists?
-          valid?
-        end
+        def exists? = valid?
+        def access_token = self.load&.fetch(:access_token, nil)
 
-        def access_token
-          tokens = self.load
-          tokens&.fetch(:access_token, nil)
-        end
+        private
 
-        def token_type
-          tokens = self.load
-          tokens&.fetch(:type, :oauth)
+        def resolve_env_key(provider)
+          default = Config::Defaults::PROVIDER_ENV_KEYS.fetch(provider, "#{provider.upcase}_API_KEY")
+          Config::Settings.new.provider_config(provider)&.fetch('env_key', nil) || default
+        rescue StandardError
+          default
         end
 
-        private
-
-        # Read Claude Code's OAuth token from macOS Keychain
         def load_from_keychain
           return nil unless RUBY_PLATFORM.include?('darwin')
 
           output = `security find-generic-password -s "#{KEYCHAIN_SERVICE}" -w 2>/dev/null`.strip
           return nil if output.empty?
 
-          data = JSON.parse(output)
-          oauth = data['claudeAiOauth']
-          return nil unless oauth && oauth['accessToken']
+          oauth = JSON.parse(output)['claudeAiOauth']
+          return nil unless oauth&.dig('accessToken')
 
-          expires_at = if oauth['expiresAt']
-                         Time.at(oauth['expiresAt'] / 1000.0) # milliseconds to seconds
-                       end
+          build_keychain_tokens(oauth)
+        rescue StandardError
+          nil
+        end
 
+        def build_keychain_tokens(oauth)
           {
             access_token: oauth['accessToken'],
             refresh_token: oauth['refreshToken'],
-            expires_at: expires_at,
+            expires_at: oauth['expiresAt'] ? Time.at(oauth['expiresAt'] / 1000.0) : nil,
             type: :oauth,
             source: :keychain
           }
-        rescue JSON::ParserError, StandardError
-          nil
         end
 
-        # Read from local YAML token file
         def load_from_file
           return nil unless File.exist?(tokens_path)
 
@@ -114,28 +130,40 @@ module RubynCode
           nil
         end
 
-        # Fall back to ANTHROPIC_API_KEY environment variable
         def load_from_env
           api_key = ENV.fetch('ANTHROPIC_API_KEY', nil)
           return nil unless api_key && !api_key.empty?
 
-          {
-            access_token: api_key,
-            refresh_token: nil,
-            expires_at: nil,
-            type: :api_key,
-            source: :env
-          }
+          { access_token: api_key, refresh_token: nil, expires_at: nil, type: :api_key, source: :env }
         end
 
-        def tokens_path
-          Config::Defaults::TOKENS_FILE
+        def write_tokens_file(data)
+          File.write(tokens_path, YAML.dump(data))
+          File.chmod(0o600, tokens_path)
         end
 
+        # Auto-encrypt a plaintext key from a pre-encryption install.
+        def migrate_plaintext_key!(data, provider, plaintext)
+          data['provider_keys'][provider.to_s] = KeyEncryption.encrypt(plaintext)
+          write_tokens_file(data)
+        rescue StandardError
+          nil # don't break reads if migration fails
+        end
+
+        def load_tokens_file
+          return nil unless File.exist?(tokens_path)
+
+          data = YAML.safe_load_file(tokens_path, permitted_classes: [Time])
+          data.is_a?(Hash) ? data : nil
+        rescue Psych::SyntaxError, Errno::EACCES
+          nil
+        end
+
+        def tokens_path = Config::Defaults::TOKENS_FILE
+
         def ensure_directory!
-          dir = File.dirname(tokens_path)
-          FileUtils.mkdir_p(dir)
-          File.chmod(0o700, dir)
+          FileUtils.mkdir_p(File.dirname(tokens_path))
+          File.chmod(0o700, File.dirname(tokens_path))
         end
 
         def parse_time(value)

data/lib/rubyn_code/autonomous/daemon.rb

@@ -14,8 +14,9 @@ module RubynCode
     #
     # Unlike the REPL, the daemon runs a full Agent::Loop per task — meaning
     # it can read files, write code, run specs, and use every tool available.
-    class Daemon
+    class Daemon # rubocop:disable Metrics/ClassLength -- daemon lifecycle + retry + audit + cost
       LIFECYCLE_STATES = %i[spawned working idle shutting_down stopped].freeze
+      MAX_TASK_RETRIES = 3
 
       attr_reader :agent_name, :role, :state, :runs_completed, :total_cost
 
@@ -32,29 +33,17 @@ module RubynCode
       # @param on_state_change [Proc, nil] callback invoked with (old_state, new_state)
       # @param on_task_complete [Proc, nil] callback invoked with (task, result_text)
       # @param on_task_error [Proc, nil] callback invoked with (task, error)
+      # @param session_persistence [Memory::SessionPersistence, nil] optional audit trail persistence
       def initialize( # rubocop:disable Metrics/ParameterLists
         agent_name:, role:, llm_client:, project_root:, task_manager:, mailbox:,
         max_runs: 100, max_cost: 10.0, poll_interval: 5, idle_timeout: 60,
-        on_state_change: nil, on_task_complete: nil, on_task_error: nil
+        on_state_change: nil, on_task_complete: nil, on_task_error: nil,
+        session_persistence: nil
       )
-        @agent_name      = agent_name
-        @role            = role
-        @llm_client      = llm_client
-        @project_root    = File.expand_path(project_root)
-        @task_manager    = task_manager
-        @mailbox         = mailbox
-        @max_runs        = max_runs
-        @max_cost        = max_cost
-        @poll_interval   = poll_interval
-        @idle_timeout    = idle_timeout
-        @on_state_change = on_state_change
-        @on_task_complete = on_task_complete
-        @on_task_error   = on_task_error
-
-        @state           = :spawned
-        @runs_completed  = 0
-        @total_cost      = 0.0
-        @stop_requested  = false
+        assign_core_attrs(agent_name:, role:, llm_client:, project_root:, task_manager:, mailbox:)
+        assign_limits(max_runs:, max_cost:, poll_interval:, idle_timeout:)
+        assign_callbacks_and_state(on_state_change, on_task_complete, on_task_error)
+        @session_persistence = session_persistence
       end
 
       # Enters the work-idle-work cycle. Blocks the calling thread until
@@ -121,6 +110,32 @@ module RubynCode
 
       # ── Signal handling ──────────────────────────────────────────
 
+      def assign_core_attrs(agent_name:, role:, llm_client:, project_root:, task_manager:, mailbox:) # rubocop:disable Metrics/ParameterLists -- mirrors constructor keyword args
+        @agent_name   = agent_name
+        @role         = role
+        @llm_client   = llm_client
+        @project_root = File.expand_path(project_root)
+        @task_manager = task_manager
+        @mailbox      = mailbox
+      end
+
+      def assign_limits(max_runs:, max_cost:, poll_interval:, idle_timeout:)
+        @max_runs      = max_runs
+        @max_cost      = max_cost
+        @poll_interval = poll_interval
+        @idle_timeout  = idle_timeout
+      end
+
+      def assign_callbacks_and_state(on_state_change, on_task_complete, on_task_error)
+        @on_state_change  = on_state_change
+        @on_task_complete = on_task_complete
+        @on_task_error    = on_task_error
+        @state            = :spawned
+        @runs_completed   = 0
+        @total_cost       = 0.0
+        @stop_requested   = false
+      end
+
       def install_signal_handlers!
         %w[INT TERM].each do |sig|
           Signal.trap(sig) { stop! }
@@ -142,16 +157,18 @@ module RubynCode
         agent_loop = build_agent_loop
         result_text = agent_loop.send_message(build_work_prompt(task))
 
-        # Accumulate cost from the budget enforcer
-        track_cost_from_enforcer(agent_loop)
+        # Accumulate cost via CostCalculator using actual token counts
+        track_cost_from_context_manager(agent_loop)
 
         # Mark the task as completed with the agent's result.
         @task_manager.complete(task.id, result: result_text)
+
+        # Persist conversation as an audit trail
+        persist_session_audit(task, agent_loop)
+
         @on_task_complete&.call(task, result_text)
       rescue StandardError => e
-        # On failure, release the task so another agent (or retry) can pick it up.
-        @task_manager.update(task.id, status: 'pending', owner: nil, result: "Error: #{e.message}")
-        @on_task_error&.call(task, e)
+        handle_task_error(task, e)
       end
 
       # Builds a fresh Agent::Loop wired with all the real tools.
@@ -181,19 +198,116 @@ module RubynCode
         )
       end
 
-      # Accumulates cost tracked by the Agent::Loop's context manager.
+      # Computes USD cost from the context manager's token counts using
+      # Observability::CostCalculator. The old approach checked for a
+      # `total_cost` method that never existed on Context::Manager, so
+      # @total_cost was always 0.0 and the max_cost safety limit never fired.
       #
       # @param agent_loop [Agent::Loop]
       # @return [void]
-      def track_cost_from_enforcer(agent_loop)
-        # The context manager tracks token usage; we extract cost if available.
-        # This is best-effort — the daemon's own total_cost is an approximation.
+      def track_cost_from_context_manager(agent_loop)
         cm = agent_loop.instance_variable_get(:@context_manager)
-        return unless cm.respond_to?(:total_cost)
+        return unless cm
+
+        tokens = extract_token_counts(cm)
+        return if tokens.values.all?(&:zero?)
+
+        model = @llm_client.respond_to?(:model) ? @llm_client.model : 'claude-sonnet-4-6'
+        @total_cost += Observability::CostCalculator.calculate(model: model, **tokens)
+      rescue StandardError
+        # Non-critical — cost tracking is best-effort
+      end
+
+      # @param context_mgr [Context::Manager]
+      # @return [Hash] :input_tokens, :output_tokens
+      def extract_token_counts(context_mgr)
+        {
+          input_tokens: context_mgr.respond_to?(:total_input_tokens) ? context_mgr.total_input_tokens.to_i : 0,
+          output_tokens: context_mgr.respond_to?(:total_output_tokens) ? context_mgr.total_output_tokens.to_i : 0
+        }
+      end
+
+      # Handles a task error with retry backoff. Increments the retry count
+      # in the task's metadata. After MAX_TASK_RETRIES, marks the task as
+      # failed instead of releasing it back to pending.
+      #
+      # @param task [Tasks::Task]
+      # @param error [StandardError]
+      # @return [void]
+      def handle_task_error(task, error)
+        retry_count = extract_retry_count(task) + 1
+
+        metadata = build_retry_metadata(task, retry_count)
+        if retry_count >= MAX_TASK_RETRIES
+          @task_manager.update(
+            task.id,
+            status: 'failed',
+            owner: nil,
+            result: "Failed after #{retry_count} retries. Last error: #{error.message}",
+            metadata: JSON.generate(metadata)
+          )
+        else
+          @task_manager.update(
+            task.id,
+            status: 'pending',
+            owner: nil,
+            result: "Error (retry #{retry_count}/#{MAX_TASK_RETRIES}): #{error.message}",
+            metadata: JSON.generate(metadata)
+          )
+        end
+        @on_task_error&.call(task, error)
+      end
+
+      # @param task [Tasks::Task]
+      # @return [Integer]
+      def extract_retry_count(task)
+        meta = parse_task_metadata(task)
+        (meta[:retry_count] || meta['retry_count'] || 0).to_i
+      end
+
+      # @param task [Tasks::Task]
+      # @param retry_count [Integer]
+      # @return [Hash]
+      def build_retry_metadata(task, retry_count)
+        meta = parse_task_metadata(task)
+        meta.merge(retry_count: retry_count)
+      end
+
+      # @param task [Tasks::Task]
+      # @return [Hash]
+      def parse_task_metadata(task)
+        raw = task.metadata
+        case raw
+        when Hash then raw
+        when String then JSON.parse(raw, symbolize_names: true)
+        else {}
+        end
+      rescue JSON::ParserError
+        {}
+      end
 
-        @total_cost += cm.total_cost.to_f
+      # Persists the agent's conversation as a session audit trail after
+      # completing a task, so there's a record of what the daemon did.
+      #
+      # @param task [Tasks::Task]
+      # @param agent_loop [Agent::Loop]
+      # @return [void]
+      def persist_session_audit(task, agent_loop)
+        return unless @session_persistence
+
+        conversation = agent_loop.instance_variable_get(:@conversation)
+        return unless conversation.respond_to?(:messages)
+
+        session_id = "daemon-#{@agent_name}-#{task.id}"
+        @session_persistence.save_session(
+          session_id: session_id,
+          project_path: @project_root,
+          messages: conversation.messages,
+          title: "Daemon: #{task.title}",
+          metadata: { agent_name: @agent_name, task_id: task.id, task_title: task.title }
+        )
       rescue StandardError
-        # Non-critical
+        # Non-critical — audit persistence is best-effort
       end
 
       # ── Idle phase ───────────────────────────────────────────────

data/lib/rubyn_code/autonomous/idle_poller.rb

@@ -35,9 +35,9 @@ module RubynCode
           return :shutdown if monotonic_now >= deadline
 
           # Messages always take priority over tasks.
-          return :resume if has_pending_messages?
+          return :resume if pending_messages?
 
-          return :resume if has_claimable_task?
+          return :resume if claimable_task?
 
           remaining = deadline - monotonic_now
           return :shutdown if remaining <= 0
@@ -53,30 +53,10 @@ module RubynCode
         @interrupted = true
       end
 
-      # Re-injects the agent's identity message when the conversation
-      # context has been compressed (i.e. the messages array is very short).
-      # This ensures the agent still knows who it is after compaction.
-      #
-      # @param messages [Array<Hash>] the current conversation messages
-      # @param identity [String] the identity/system prompt to re-inject
-      # @param threshold [Integer] message count below which re-injection triggers (default 3)
-      # @return [void]
-      def self.reinject_identity(messages, identity:, threshold: 3)
-        return if messages.length >= threshold
-        return if identity.nil? || identity.empty?
-
-        # Only re-inject if the identity is not already present as the
-        # first user message.
-        first_user = messages.find { |m| m[:role] == 'user' }
-        return if first_user && first_user[:content].to_s.include?(identity[0, 100])
-
-        messages.unshift({ role: 'user', content: identity })
-      end
-
       private
 
       # @return [Boolean]
-      def has_pending_messages?
+      def pending_messages?
         messages = @mailbox.pending_for(@agent_name)
         messages.is_a?(Array) ? !messages.empty? : false
       rescue StandardError
@@ -84,7 +64,7 @@ module RubynCode
       end
 
       # @return [Boolean]
-      def has_claimable_task?
+      def claimable_task?
         rows = @task_manager.db.query(<<~SQL).to_a
           SELECT 1 FROM tasks
           WHERE status = 'pending'

data/lib/rubyn_code/autonomous/task_claimer.rb

@@ -6,56 +6,21 @@ module RubynCode
     # Uses optimistic locking to handle race conditions when multiple
     # agents attempt to claim the same task concurrently.
     module TaskClaimer
-      # Finds the first ready (pending, unowned) task, claims it for the
-      # given agent, and returns the updated Task. Returns nil if no work
-      # is available.
+      MAX_RETRIES = 3
+
+      # Finds the first ready (pending, unowned) task that hasn't exceeded
+      # max retries, claims it for the given agent, and returns the updated
+      # Task. Returns nil if no work is available.
       #
       # @param task_manager [#db, #update_task, #list_tasks] task persistence layer
       # @param agent_name [String] unique identifier of the claiming agent
+      # @param max_retries [Integer] maximum retry count before skipping a task
       # @return [Tasks::Task, nil] the claimed task, or nil if none available
-      def self.call(task_manager:, agent_name:)
+      def self.call(task_manager:, agent_name:, max_retries: MAX_RETRIES)
         db = task_manager.db
-
-        # Atomically claim the first eligible task. The WHERE conditions
-        # ensure that only pending tasks with no current owner are touched,
-        # avoiding race conditions with other agents.
-        db.execute(<<~SQL, [agent_name])
-          UPDATE tasks
-          SET owner = ?,
-              status = 'in_progress',
-              updated_at = datetime('now')
-          WHERE id = (
-            SELECT id FROM tasks
-            WHERE status = 'pending'
-              AND (owner IS NULL OR owner = '')
-            ORDER BY priority DESC, created_at ASC
-            LIMIT 1
-          )
-          AND status = 'pending'
-          AND (owner IS NULL OR owner = '')
-        SQL
-
-        # Fetch the task we just claimed. Using owner + status filters
-        # ensures we only retrieve a task that *this* agent successfully
-        # claimed (another agent cannot have flipped it in between).
-        rows = db.query(<<~SQL, [agent_name]).to_a
-          SELECT id, session_id, title, description, status,
-                 priority, owner, result, metadata, created_at, updated_at
-          FROM tasks
-          WHERE owner = ?
-            AND status = 'in_progress'
-          ORDER BY updated_at DESC
-          LIMIT 1
-        SQL
-
-        return nil if rows.empty?
-
-        row = rows.first
-        build_task(row)
+        claim_next_pending_task(db, agent_name, max_retries)
+        fetch_claimed_task(db, agent_name)
       rescue StandardError => e
-        # If anything goes wrong (e.g. task was already claimed between
-        # our SELECT and UPDATE, or a constraint violation) we treat it
-        # as "no work available" rather than crashing the daemon.
         RubynCode.logger.warn("TaskClaimer: failed to claim task: #{e.message}") if RubynCode.respond_to?(:logger)
         nil
       end
@@ -63,6 +28,43 @@ module RubynCode
       class << self
         private
 
+        def claim_next_pending_task(db, agent_name, max_retries)
+          db.execute(<<~SQL, [agent_name, max_retries])
+            UPDATE tasks
+            SET owner = ?,
+                status = 'in_progress',
+                updated_at = datetime('now')
+            WHERE id = (
+              SELECT t.id FROM tasks t
+              WHERE t.status = 'pending'
+                AND (t.owner IS NULL OR t.owner = '')
+                AND COALESCE(
+                  json_extract(t.metadata, '$.retry_count'), 0
+                ) < ?
+              ORDER BY t.priority DESC, t.created_at ASC
+              LIMIT 1
+            )
+            AND status = 'pending'
+            AND (owner IS NULL OR owner = '')
+          SQL
+        end
+
+        def fetch_claimed_task(db, agent_name)
+          rows = db.query(<<~SQL, [agent_name]).to_a
+            SELECT id, session_id, title, description, status,
+                   priority, owner, result, metadata, created_at, updated_at
+            FROM tasks
+            WHERE owner = ?
+              AND status = 'in_progress'
+            ORDER BY updated_at DESC
+            LIMIT 1
+          SQL
+
+          return nil if rows.empty?
+
+          build_task(rows.first)
+        end
+
         # @param row [Hash] a database row hash
         # @return [Tasks::Task]
         def build_task(row)