Move agent providers into shared registry

This will allow us to support agents other than Claude Code

Author: dersam

lib/roast/cogs/agent.rb

@@ -48,26 +48,14 @@ class MissingPromptError < AgentCogError; end
       #: (Input) -> Output
       def execute(input)
         puts "[USER PROMPT] #{input.valid_prompt!}" if config.show_prompt?
-        output = provider.invoke(input)
+        output = config.values[:provider].invoke(input)
         # NOTE: If progress is displayed, the agent's response will always be the last progress message,
         # so showing it again is duplicative.
         puts "[AGENT RESPONSE] #{output.response}" if config.show_response? && !config.show_progress?
         puts "[AGENT STATS] #{output.stats}" if config.show_stats?
         puts "Session ID: #{output.session}" if config.show_stats?
         output
       end
-
-      private
-
-      #: () -> Provider
-      def provider
-        @provider ||= case config.valid_provider!
-        when :claude
-          Providers::Claude.new(config)
-        else
-          raise UnknownProviderError, "Unknown provider: #{config.valid_provider!}"
-        end
-      end
     end
   end
 end

lib/roast/cogs/agent/config.rb

@@ -5,7 +5,17 @@ module Roast
   module Cogs
     class Agent < Cog
       class Config < Cog::Config
-        VALID_PROVIDERS = [:claude].freeze #: Array[Symbol]
+        #: (?Hash[Symbol, untyped]) -> void
+        def initialize(initial = {})
+          super(initial)
+
+          @provider_registry = {}
+        end
+
+        def validate!
+          # Provider registry is responsible for ensuring the validity of providers
+          @values[:provider] = @provider_registry.fetch(@values[:provider]).new(self)
+        end
 
         # Configure the cog to use a specified provider when invoking an agent
         #
@@ -36,27 +46,6 @@ def use_default_provider!
           @values[:provider] = nil
         end
 
-        # Get the validated provider name that the cog is configured to use when invoking an agent
-        #
-        # Note: this method will return the name of a valid provider or raise an `InvalidConfigError`.
-        # It will __not__, however, validate that the agent is properly installed on your system.
-        # If the agent is not properly installed, you will likely experience a failure when Roast attempts to
-        # run your workflow.
-        #
-        # #### See Also
-        # - `provider`
-        # - `use_default_provider!`
-        #
-        #: () -> Symbol
-        def valid_provider!
-          provider = @values[:provider] || VALID_PROVIDERS.first
-          unless VALID_PROVIDERS.include?(provider)
-            raise ArgumentError, "'#{provider}' is not a valid provider. Available providers include: #{VALID_PROVIDERS.join(", ")}"
-          end
-
-          provider
-        end
-
         # Configure the cog to use a specific base command when invoking the agent
         #
         # The command format is provider-specific.

lib/roast/config_manager.rb

@@ -8,9 +8,10 @@ class ConfigManagerNotPreparedError < ConfigManagerError; end
     class ConfigManagerAlreadyPreparedError < ConfigManagerError; end
     class IllegalCogNameError < ConfigManagerError; end
 
-    #: (Cog::Registry, Array[^() -> void]) -> void
-    def initialize(cog_registry, config_procs)
+    #: (Cog::Registry, Array[^() -> void], ProviderRegistry) -> void
+    def initialize(cog_registry, config_procs, provider_registry)
       @cog_registry = cog_registry
+      @provider_registry = provider_registry
       @config_procs = config_procs
       @config_context = ConfigContext.new #: ConfigContext
       @global_config = Cog::Config.new #: Cog::Config
@@ -54,6 +55,16 @@ def config_for(cog_class, name = nil)
         name_scoped_config = fetch_name_scoped_config(cog_class, name)
         config = config.merge(name_scoped_config)
       end
+
+      # Special case for agent cog. Insert the provider registry.
+      # This is the only cog that needs this right now - revisit in the future to
+      # see if we need a way to define hooks for custom cog data to be distributed
+      # at config time.
+      # NOTE: This must happen after all merges, since merge creates new Config instances.
+      if config.is_a?(Cogs::Agent::Config)
+        config.instance_variable_set(:@provider_registry, @provider_registry)
+      end
+
       config.validate!
       config
     end

lib/roast/provider_registry.rb

@@ -0,0 +1,45 @@
+# typed: true
+# frozen_string_literal: true
+
+module Roast
+  # Maintains the list of registered agent providers
+  # Built in providers are registered automatically.
+  # Custom agents are registered per workflow with `use`.)
+  class ProviderRegistry
+    class ProviderRegistryError < Roast::Error; end
+    class DuplicateProviderNameError < ProviderRegistryError; end
+    class ProviderNotFoundError < ProviderRegistryError; end
+
+    #: Symbol
+    attr_accessor :default
+
+    delegate :key?, to: :@providers
+
+    def initialize
+      @providers = {} #: Hash[Symbol, singleton(Cogs::Agent::Provider)]
+      @default = ENV["ROAST_DEFAULT_AGENT"]&.to_sym || :claude
+    end
+
+    #: (singleton(Cogs::Agent::Provider), ?Symbol?) -> void
+    def register(provider_class, name = nil)
+      name = build_provider_name(provider_class) if name.blank?
+      raise DuplicateProviderNameError if @providers.key?(name)
+
+      @providers[name] = provider_class
+    end
+
+    def fetch(name)
+      name = default if name.nil?
+      raise ProviderNotFoundError unless @providers.key?(name)
+
+      @providers.fetch(name)
+    end
+
+    private
+
+    #: (singleton(Cogs::Agent::Provider)) -> Symbol
+    def build_provider_name(provider_class)
+      provider_class.name.not_nil!.demodulize.underscore.to_sym
+    end
+  end
+end

lib/roast/workflow.rb

@@ -28,6 +28,7 @@ def initialize(workflow_path, workflow_context)
       @workflow_context = workflow_context #: WorkflowContext
       @workflow_definition = File.read(workflow_path) #: String
       @cog_registry = Cog::Registry.new #: Cog::Registry
+      @provider_registry = ProviderRegistry.new #: ProviderRegistry
       @config_procs = [] #: Array[^() -> void]
       @execution_procs = { nil: [] } #: Hash[Symbol?, Array[^() -> void]]
       @config_manager = nil #: ConfigManager?
@@ -40,7 +41,8 @@ def prepare!
 
       @preparing = true
       extract_dsl_procs!
-      @config_manager = ConfigManager.new(@cog_registry, @config_procs)
+      add_providers!
+      @config_manager = ConfigManager.new(@cog_registry, @config_procs, @provider_registry)
       @config_manager.not_nil!.prepare!
       # TODO: probably we should just not pass the params as the top-level scope value anymore
       @execution_manager = ExecutionManager.new(@cog_registry, @config_manager.not_nil!, @execution_procs, @workflow_context, scope_value: @workflow_context.params)
@@ -118,5 +120,9 @@ def use(cogs = [], from: nil)
     def extract_dsl_procs!
       instance_eval(@workflow_definition, @workflow_path.realpath.to_s, 1)
     end
+
+    def add_providers!
+      @provider_registry.register(Roast::Cogs::Agent::Providers::Claude, :claude)
+    end
   end
 end

test/roast/cogs/agent/config_test.rb

@@ -8,35 +8,63 @@ class Agent < Cog
       class ConfigTest < ActiveSupport::TestCase
         def setup
           @config = Config.new
-          @default_provider = Config::VALID_PROVIDERS.first
+          @default_provider = :claude
         end
 
         # Provider configuration tests
         test "provider sets provider value" do
           @config.provider(@default_provider)
 
-          assert_equal @default_provider, @config.valid_provider!
+          assert_equal @default_provider, @config.values[:provider]
         end
 
         test "use_default_provider! clears provider value" do
           @config.provider(:fake_provider)
           @config.use_default_provider!
 
-          assert_equal @default_provider, @config.valid_provider!
+          assert_nil @config.values[:provider]
         end
 
-        test "valid_provider! returns default when not set" do
-          assert_equal @default_provider, @config.valid_provider!
+        test "initialize sets provider_registry to empty hash" do
+          config = Config.new
+
+          assert_equal({}, config.instance_variable_get(:@provider_registry))
         end
 
-        test "valid_provider! raises on invalid provider" do
-          @config.provider(:invalid_provider)
+        test "validate! resolves provider from registry and instantiates with config" do
+          mock_provider_instance = mock("provider_instance")
+          mock_provider_class = mock("provider_class")
+          mock_provider_class.expects(:new).with(@config).returns(mock_provider_instance)
 
-          error = assert_raises(ArgumentError) do
-            @config.valid_provider!
-          end
+          registry = { claude: mock_provider_class }
+          @config.provider(:claude)
+          @config.instance_variable_set(:@provider_registry, registry)
+          @config.validate!
+
+          assert_equal mock_provider_instance, @config.values[:provider]
+        end
+
+        test "validate! uses nil provider key when provider not explicitly set" do
+          mock_provider_instance = mock("provider_instance")
+          mock_provider_class = mock("provider_class")
+          mock_provider_class.expects(:new).with(@config).returns(mock_provider_instance)
 
-          assert_match(/invalid_provider.*not a valid provider/, error.message)
+          registry = mock("registry")
+          registry.expects(:fetch).with(nil).returns(mock_provider_class)
+          @config.instance_variable_set(:@provider_registry, registry)
+          @config.validate!
+
+          assert_equal mock_provider_instance, @config.values[:provider]
+        end
+
+        test "validate! raises when provider not found in registry" do
+          registry = ProviderRegistry.new
+          @config.provider(:nonexistent)
+          @config.instance_variable_set(:@provider_registry, registry)
+
+          assert_raises(ProviderRegistry::ProviderNotFoundError) do
+            @config.validate!
+          end
         end
 
         # Command configuration tests

test/roast/config_manager_test.rb

@@ -22,15 +22,15 @@ def setup
     end
 
     test "prepare! transitions to prepared state" do
-      manager = ConfigManager.new(@registry, [])
+      manager = ConfigManager.new(@registry, [], ProviderRegistry.new)
 
       refute manager.prepared?
       manager.prepare!
       assert manager.prepared?
     end
 
     test "prepare! raises when called twice" do
-      manager = ConfigManager.new(@registry, [])
+      manager = ConfigManager.new(@registry, [], ProviderRegistry.new)
       manager.prepare!
 
       assert_raises(ConfigManager::ConfigManagerAlreadyPreparedError) do
@@ -44,22 +44,22 @@ def setup
         test_cog { timeout 60 }
         timeout_set = true
       end
-      manager = ConfigManager.new(@registry, [config_proc])
+      manager = ConfigManager.new(@registry, [config_proc], ProviderRegistry.new)
       manager.prepare!
 
       assert timeout_set
     end
 
     test "config_for raises when not prepared" do
-      manager = ConfigManager.new(@registry, [])
+      manager = ConfigManager.new(@registry, [], ProviderRegistry.new)
 
       assert_raises(ConfigManager::ConfigManagerNotPreparedError) do
         manager.config_for(TestCog)
       end
     end
 
     test "config_for returns default config when no config procs are provided" do
-      manager = ConfigManager.new(@registry, [])
+      manager = ConfigManager.new(@registry, [], ProviderRegistry.new)
       manager.prepare!
 
       config = manager.config_for(TestCog)
@@ -71,7 +71,7 @@ def setup
       config_proc = proc do
         test_cog { timeout 60 }
       end
-      manager = ConfigManager.new(@registry, [config_proc])
+      manager = ConfigManager.new(@registry, [config_proc], ProviderRegistry.new)
       manager.prepare!
 
       config = manager.config_for(TestCog)
@@ -83,7 +83,7 @@ def setup
       config_proc = proc do
         test_cog(:my_step) { timeout 90 }
       end
-      manager = ConfigManager.new(@registry, [config_proc])
+      manager = ConfigManager.new(@registry, [config_proc], ProviderRegistry.new)
       manager.prepare!
 
       scoped_config = manager.config_for(TestCog, :my_step)
@@ -97,7 +97,7 @@ def setup
       config_proc = proc do
         test_cog(/^api_/) { timeout 120 }
       end
-      manager = ConfigManager.new(@registry, [config_proc])
+      manager = ConfigManager.new(@registry, [config_proc], ProviderRegistry.new)
       manager.prepare!
 
       matching_config = manager.config_for(TestCog, :api_call)
@@ -112,7 +112,7 @@ def setup
         test_cog { async! }
         test_cog(:my_step) { timeout 90 }
       end
-      manager = ConfigManager.new(@registry, [config_proc])
+      manager = ConfigManager.new(@registry, [config_proc], ProviderRegistry.new)
       manager.prepare!
 
       config = manager.config_for(TestCog, :my_step)
@@ -125,14 +125,37 @@ def setup
       config_proc = proc do
         global { abort_on_failure! }
       end
-      manager = ConfigManager.new(@registry, [config_proc])
+      manager = ConfigManager.new(@registry, [config_proc], ProviderRegistry.new)
       manager.prepare!
 
       config = manager.config_for(TestCog)
 
       assert config.abort_on_failure?
     end
 
+    test "config_for injects provider registry into Agent::Config" do
+      registry = Cog::Registry.new
+      provider_registry = ProviderRegistry.new
+      provider_registry.register(Cogs::Agent::Providers::Claude, :claude)
+
+      manager = ConfigManager.new(registry, [], provider_registry)
+      manager.prepare!
+
+      config = manager.config_for(Cogs::Agent)
+
+      assert_equal provider_registry, config.instance_variable_get(:@provider_registry)
+    end
+
+    test "config_for does not inject provider registry into non-Agent configs" do
+      manager = ConfigManager.new(@registry, [], ProviderRegistry.new)
+      manager.prepare!
+
+      config = manager.config_for(TestCog)
+
+      refute config.respond_to?(:provider_registry)
+      refute config.instance_variable_defined?(:@provider_registry)
+    end
+
     test "prepare! raises IllegalCogNameError when cog name conflicts with existing method" do
       # Register a cog whose derived name ("freeze") conflicts with Object#freeze
       conflicting_cog = Class.new(Cog) do
@@ -144,7 +167,7 @@ def name
       end
       @registry.use(conflicting_cog)
 
-      manager = ConfigManager.new(@registry, [])
+      manager = ConfigManager.new(@registry, [], ProviderRegistry.new)
 
       assert_raises(ConfigManager::IllegalCogNameError) do
         manager.prepare!

test/roast/execution_manager_test.rb

@@ -8,7 +8,7 @@ def setup
       @registry = Cog::Registry.new
       @registry.use(TestCogSupport::TestCog)
 
-      @config_manager = ConfigManager.new(@registry, [])
+      @config_manager = ConfigManager.new(@registry, [], ProviderRegistry.new)
       @config_manager.prepare!
 
       @workflow_context = WorkflowContext.new(
@@ -68,7 +68,7 @@ def name
       conflicting_registry.use(conflicting_cog)
 
       # Use a clean registry for config_manager so it doesn't hit the same conflict
-      config_manager = ConfigManager.new(Cog::Registry.new, [])
+      config_manager = ConfigManager.new(Cog::Registry.new, [], ProviderRegistry.new)
       config_manager.prepare!
 
       manager = ExecutionManager.new(