ruby_llm-code_mode 0.4.0 → 0.4.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: fb23ad192879e9ce6bcac0204c62c44ad068943dadd5e46a3a2e82bd177411c7
-  data.tar.gz: 39b52ca9cf5f733224a1a8a999ae50cbe091467819d86f63a7c38610be665233
+  metadata.gz: 71758c99ee2449675bbfb17951f4f12e6d9a4e35b7bd414cf9a5ff2a5219cddc
+  data.tar.gz: 0ebbe47c2d6b8267583edee716be3f6039cdf27c67debc699dbdf960e8c4e0b5
 SHA512:
-  metadata.gz: 6e41645fd4404b2a70e3a3606ab6b1aba047e20e1e610278ddaac8ae85b8512d0c0af134a897fface831e7a1547b77610444bbb8dffb6a80b468bf4a79f09af2
-  data.tar.gz: 60adf578f5deafbd36822e14179e8bcd74856ba6d4b598910cc2be40520ef577f8fd534ffde0ecad2d5ab5c512ff2ac2b98ae60f8c66b0cc0dfa8253bde5b609
+  metadata.gz: a95977801ea3cf1fe0c057e2e69e3f965145306e6d2edc588d66c0cc953f0f219ee22da193bed054695ad550e4d921fa29bb36c9d9eaa200a3866953fb94f667
+  data.tar.gz: 0a394a1b1c7aafa6d44240b13afb28096dc31a311d22cb6713c2847893c3e1d879c1b0a99732140f695fe65165d62110b8c3ebb982155267bac4c25eeb628b5d

data/README.md

@@ -1,4 +1,4 @@
-# ruby_llm-code_mode
+# RubyLLM::CodeMode
 
 A [RubyLLM](https://rubyllm.com) tool that runs model-generated Ruby code inside a
 secure sandbox, powered by [SecurityBox](https://rubygems.org/gems/security_box).
@@ -269,4 +269,4 @@ rake release                      # specs + git tag vX.Y.Z + push to GitHub + pu
 ```
 
 `rake release` refuses a dirty working tree or an existing tag and uses the
-RubyGems credentials from `~/.gem/credentials`.
+RubyGems credentials from `~/.gem/credentials`.

data/lib/ruby_llm/code_mode/description.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class CodeMode
+    FIXED_DESCRIPTION = <<~DESC.freeze
+      Executes Ruby code inside a secure sandbox and returns what it produced.
+
+      ## How to use
+
+      Pass a complete, self-contained Ruby script as the `code` argument. It runs
+      as a fresh Ruby 4.0 process with:
+      - no network, no threads, no subprocesses, and no access to the host
+        filesystem outside the mounted folders listed below;
+      - a private writable scratch directory at `/work`, wiped between executions;
+      - no state carried over between executions, except the contents of
+        read-write folders.
+
+      Capture progress with `puts` (returned as `stdout`) and end with a value or
+      a JSON-serializable data structure (returned as `value`; non-serializable
+      objects come back as their inspect string). If the code raises, the error
+      class, message and backtrace are returned so you can fix and retry.
+    DESC
+
+    HOST_TOOLS_INTRO = <<~DESC.freeze
+      Registered host tools are callable from inside the sandbox. Call one with
+      `SB.call('name', key: value)`: arguments must be JSON-serializable and become
+      the tool's keyword arguments, and the return value is the tool's result
+      (non-serializable values come back as their inspect string). Bad arguments
+      come back as an `error` hash; handler failures raise `SB::ToolError`
+      (rescuable); unknown names raise `SB::UnknownTool`. Limits: at most 1000
+      tool calls per execution and 1 MiB per result.
+    DESC
+
+    class << self
+      # Builds the description the model sees for this tool: the fixed
+      # "How to use" text, one section per mount kind listing each mounted
+      # folder, the "Host tools (call with SB.call)" section with every
+      # bound tool and its parameters, and the "Server notes" section with
+      # each MCP server's instructions. `extra_tools`/`extra_mcps` carry
+      # the per-instance additions (see #add_tools/#add_mcps); with no
+      # extras the output is identical to the class-level description.
+      def build_description(extra_tools: {}, extra_mcps: [])
+        read_only, read_write = mounts.partition { |m| m.mode == :read_only }
+        [
+          FIXED_DESCRIPTION.rstrip,
+          folder_section("## Read-only folders (readable, never writable)", read_only),
+          folder_section("## Read-write folders (readable and writable)", read_write),
+          Registry.tools_section(tools.merge(extra_tools)),
+          Registry.server_notes_section(mcps + extra_mcps)
+        ].compact.join("\n\n")
+      end
+
+      # Reads the tool description the model sees (the build_description
+      # output). Passing a text raises ArgumentError: the description is
+      # built from the declared mounts, bound tools and MCP servers, and
+      # cannot be set by hand.
+      def description(text = nil)
+        raise ArgumentError,
+              "RubyLLM::CodeMode builds the tool description from the declared mounts; " \
+              "the description is fixed and cannot be set by hand" if text
+
+        build_description
+      end
+
+      private
+
+      def folder_section(header, entries)
+        return nil if entries.empty?
+
+        lines = entries.map do |entry|
+          line = "- `#{entry.guest}`"
+          line += " — #{entry.description}" unless entry.description.empty?
+          line
+        end
+        [header, *lines].join("\n")
+      end
+    end
+  end
+end

data/lib/ruby_llm/code_mode/dsl.rb

@@ -0,0 +1,139 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class CodeMode
+    class << self
+      # The tool name presented to the model: the underscored class-name leaf
+      # (`MyApp::Tools::Weather` → `weather`, `SalesNotes` → `sales_notes`),
+      # without the `_tool` suffix. The base RubyLLM::Tool.tool_name would
+      # leak the namespace (`ruby_llm--code_mode`), so CodeMode derives the
+      # name from the leaf only.
+      def tool_name
+        leaf = name.to_s.split('::').last
+        return '' unless leaf
+
+        RubyLLM::Support::Utils.underscore(leaf).delete_suffix('_tool')
+      end
+
+      # Mounts a host folder into the sandbox as read-only. `source:` is
+      # expanded against the boot working directory (absolute paths pass
+      # through); `dest:` must be absolute, unique and outside the reserved
+      # trees (`/work`, `/usr`, `/src`). Invalid declarations raise
+      # ArgumentError at class-definition time, so a misconfigured tool
+      # never reaches a live chat.
+      def mount(source:, dest:, description: nil)
+        add_mount(source, dest, :read_only, description)
+      end
+
+      # Same as .mount, but the sandboxed code can create, modify and delete
+      # files inside the folder (enforced by wasmtime; the contents persist
+      # between executions of the same tool instance).
+      def mount_rw(source:, dest:, description: nil)
+        add_mount(source, dest, :read_write, description)
+      end
+
+      # Binds RubyLLM::Tool classes/instances as host tools callable from
+      # inside the sandbox via SB.call:
+      #
+      #   tools Weather                     # class (instantiated with .new)
+      #   tools Weather, Echo               # several at once
+      #   tools [Weather, Echo]             # or arrays
+      #   tools "forecast" => Weather       # explicit sandbox name
+      #   tools forecast: Weather           # same, as kwargs
+      #
+      # Without arguments it returns the hash of bound tools. The name is
+      # derived from the class-name leaf (`MyApp::Tools::Weather` binds as
+      # `weather`). Fails fast: every item must be a RubyLLM::Tool (class or
+      # instance), names must be unique and non-empty, and if anything goes
+      # wrong nothing is bound.
+      def tools(*args, **kwargs)
+        return @tools ||= {} if args.empty? && kwargs.empty?
+
+        Registry.register_tools(tools, args, kwargs)
+      end
+
+      # Connects MCP server(s) and binds all their tools into the sandbox:
+      #
+      #   docs = RubyLLM.mcp(url: "https://learn.microsoft.com/api/mcp")
+      #
+      #   class DocsResearch < RubyLLM::CodeMode
+      #     mcps docs                    # instance
+      #     mcps MicrosoftDocs           # or the class (instantiated with .new)
+      #     mcps Linear, user: current_user  # class + declared inputs
+      #     mcps Docs, Github            # several servers / arrays at once
+      #   end
+      #
+      # Without arguments it returns the array of connected servers. Server
+      # tools keep their names (`microsoft_docs_search`, ...); a tool whose
+      # name is already bound is re-bound as "<server>_tool". The servers'
+      # instructions (when the server sends them) are snapshotted at
+      # definition and listed in the description under "## Server notes".
+      # Fails fast: every argument must be a RubyLLM::MCP (class or
+      # instance), and if anything goes wrong — a server that cannot be
+      # reached, a name that stays colliding after prefixing — nothing is
+      # bound.
+      def mcps(*servers, **inputs)
+        return @mcps ||= [] if servers.empty? && inputs.empty?
+
+        Registry.register_mcps(tools, mcps, servers, inputs)
+      end
+
+      # The mounts declared so far, as Mount structs (read-only and
+      # read-write together, in declaration order). Subclasses start with a
+      # copy of the parent's list, so adding to a child never touches the
+      # parent.
+      def mounts
+        @mounts ||= []
+      end
+
+      # Builds and memoizes the SecurityBox::Configuration for this class:
+      # the default limits (timeout_ms: 30_000, fuel_ms: 10_000), the
+      # declared mounts and one RPC handler per bound tool. Memoized per
+      # class; subclasses build their own. Tool instances without
+      # per-instance additions share it (see #add_tools).
+      def configuration
+        @configuration ||= SecurityBox::Configuration.build(
+          timeout_ms: DEFAULT_TIMEOUT_MS,
+          fuel_ms: DEFAULT_FUEL_MS,
+          mounts: mounts.map(&:payload),
+          rpcs: Registry.rpc_handlers(tools)
+        )
+      end
+
+      # Pays the WebAssembly compilation cost up front (about 15 seconds on
+      # the first run, cached afterwards in ~/.cache/security_box/modules),
+      # so the first real execution takes a few hundred milliseconds
+      # instead. Call it once at boot: Analytics.warmup.
+      def warmup
+        SecurityBox.warmup(timeout_ms: DEFAULT_TIMEOUT_MS, fuel_ms: DEFAULT_FUEL_MS)
+      end
+
+      # Ruby hook: subclasses start with a copy of the parent's mounts,
+      # tools and MCP servers (entries and instances shared), and their own
+      # configuration memo so limits/rpcs are rebuilt per class.
+      def inherited(subclass)
+        super
+        subclass.instance_variable_set(:@mounts, mounts.dup)
+        subclass.instance_variable_set(:@tools, tools.dup)
+        subclass.instance_variable_set(:@mcps, mcps.dup)
+        subclass.instance_variable_set(:@configuration, nil)
+      end
+
+      private
+
+      def add_mount(source, dest, mode, description)
+        mounts << Mount.new(
+          host: File.expand_path(source.to_s),
+          guest: dest.to_s,
+          mode: mode,
+          description: description.to_s
+        )
+        SecurityBox::Mounts.normalize(mounts.map(&:payload))
+      rescue SecurityBox::InvalidConfiguration => e
+        mounts.pop
+        raise ArgumentError,
+              "Invalid mount #{source.inspect} => #{dest.inspect}: #{e.message}"
+      end
+    end
+  end
+end

data/lib/ruby_llm/code_mode/instance_methods.rb

@@ -0,0 +1,109 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class CodeMode
+    parameter :code, description: "Complete Ruby script to execute in the sandbox"
+
+    # Binds more host tools onto this specific instance, on top of the
+    # class-level ones. Accepts the same forms as the class-level .tools:
+    # classes, ready instances, several at once, arrays, and
+    # "name" => tool pairs (hash or kwargs). Instances of the same class are
+    # unaffected, and the additions show up in this instance's description,
+    # so chats built after the call expose them to the model. Fails fast and
+    # rolls everything back; once the first execution happened the sandbox is
+    # already built and further additions raise.
+    def add_tools(*args, **kwargs)
+      ensure_extendable!
+
+      Registry.register_tools(instance_tools, args, kwargs, reserved: self.class.tools)
+    end
+
+    # Connects MCP server(s) onto this specific instance, on top of the
+    # class-level ones. Same forms and fail-fast semantics as .mcps
+    # (instances, classes with their declared inputs, arrays); the servers'
+    # tools bind into this instance only. Raises after the first execution.
+    def add_mcps(*servers, **inputs)
+      ensure_extendable!
+
+      Registry.register_mcps(instance_tools, instance_mcps, servers, inputs, reserved: self.class.tools)
+    end
+
+    # The description this instance presents to the model: the class-level
+    # build plus the tools and MCP servers added with #add_tools/#add_mcps.
+    # RubyLLM reads it when the tool is registered with a chat, so call
+    # chat.with_tools after the additions.
+    def description
+      self.class.build_description(extra_tools: instance_tools, extra_mcps: instance_mcps)
+    end
+
+    # Runs the model-generated Ruby script inside the sandbox and returns
+    # the result as a hash for the model: `status` ("ok", "error",
+    # "timeout", "fuel_exhausted", "memory_limit" or "sandbox_error"),
+    # `value` on success, `error` otherwise, plus `stdout`/`stderr` when
+    # the guest produced output. Sandbox failures (missing image, invalid
+    # configuration) come back as a recoverable "sandbox_error" instead of
+    # raising.
+    def execute(code:)
+      format_result(sandbox.eval(code))
+    rescue SecurityBox::Error => e
+      { status: "sandbox_error", error: { class: e.class.name, message: e.message } }
+    end
+
+    private
+
+    def ensure_extendable!
+      return unless instance_variable_defined?(:@sandbox)
+
+      raise ArgumentError,
+            "cannot add tools or MCP servers after the first execution; " \
+            "create a new #{self.class.name} instance instead"
+    end
+
+    def instance_tools
+      @instance_tools ||= {}
+    end
+
+    def instance_mcps
+      @instance_mcps ||= []
+    end
+
+    # Instances without additions share the class configuration; instances
+    # with additions build their own, merging the class rpcs with the
+    # instance ones (names cannot collide — the registrations validate that).
+    def configuration
+      return self.class.configuration if instance_tools.empty? && instance_mcps.empty?
+
+      @configuration ||= SecurityBox::Configuration.build(
+        timeout_ms: DEFAULT_TIMEOUT_MS,
+        fuel_ms: DEFAULT_FUEL_MS,
+        mounts: self.class.mounts.map(&:payload),
+        rpcs: self.class.configuration.rpcs.merge(Registry.rpc_handlers(instance_tools))
+      )
+    end
+
+    def sandbox
+      @sandbox ||= SecurityBox::Sandbox.new(configuration)
+    end
+
+    def format_result(result)
+      payload = { status: result.status.to_s }
+      case result.status
+      when :ok
+        payload[:value] = result.value
+      when :error
+        payload[:error] = result.error || { "class" => "Unknown", "message" => "guest code failed" }
+      when :timeout
+        payload[:error] = "execution timed out before finishing"
+      when :fuel_exhausted
+        payload[:error] = "execution exceeded its CPU budget before finishing"
+      when :memory_limit
+        payload[:error] = "execution exceeded the sandbox memory limit"
+      when :sandbox_error
+        payload[:error] = "the sandbox failed to run the code"
+      end
+      payload[:stdout] = result.stdout if result.stdout && !result.stdout.empty?
+      payload[:stderr] = result.stderr if result.stderr && !result.stderr.empty?
+      payload
+    end
+  end
+end

data/lib/ruby_llm/code_mode/registry.rb

@@ -0,0 +1,233 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class CodeMode
+    # Shared registration and rendering core, parameterized by the target
+    # collections: the class-level DSL (tools/mcps) and the instance-level API
+    # (add_tools/add_mcps) run the exact same validation and rollback logic,
+    # only over different hashes/arrays.
+    module Registry
+      EMPTY = {}.freeze
+
+      class << self
+        # Binds tools into `target` (a name => ToolEntry hash). Accepts tool
+        # classes/instances, arrays of them, and "name" => tool pairs (hash
+        # arguments or kwargs). All-or-nothing: on any failure `target` is
+        # restored and the error re-raised. `reserved` lists names already
+        # taken elsewhere (e.g. the class-level bindings for instance
+        # additions); duplicates against it raise.
+        def register_tools(target, args, kwargs, reserved: EMPTY)
+          snapshot = target.dup
+          begin
+            args.flatten(1).each do |item|
+              if item.is_a?(Hash)
+                item.each { |name, bound| add_tool(target, bound, name.to_s, reserved) }
+              else
+                add_tool(target, item, nil, reserved)
+              end
+            end
+            kwargs.each { |name, bound| add_tool(target, bound, name.to_s, reserved) }
+          rescue StandardError
+            target.replace(snapshot)
+            raise
+          end
+          target
+        end
+
+        # Connects MCP server(s) and binds all their tools into
+        # `target_tools`/`target_mcps`. All-or-nothing across both
+        # collections.
+        def register_mcps(target_tools, target_mcps, servers, inputs, reserved: EMPTY)
+          snapshot_tools = target_tools.dup
+          snapshot_mcps = target_mcps.dup
+          flat = servers.flatten(1)
+          raise ArgumentError, "expected at least one RubyLLM::MCP server" if flat.empty?
+          unless inputs.empty? || flat.all?(Class)
+            raise ArgumentError, "inputs only apply when passing RubyLLM::MCP classes"
+          end
+
+          normalize_mcp_servers(flat, inputs).each do |server|
+            bind_mcp_server(target_tools, target_mcps, server, reserved)
+          end
+          target_mcps
+        rescue StandardError
+          target_tools.replace(snapshot_tools)
+          target_mcps.replace(snapshot_mcps)
+          raise
+        end
+
+        # One RPC handler per bound tool, ready for
+        # SecurityBox::Configuration.build(rpcs:). Each handler validates
+        # that the guest passed a hash of arguments, forwards the call to
+        # the tool instance (string keys symbolized, schema mistakes coming
+        # back as RubyLLM's recoverable error hash) and normalizes MCP
+        # Result objects into plain data for the guest.
+        def rpc_handlers(tools)
+          tools.transform_values { |entry| ->(args) { call_bound_tool(entry, args) } }
+        end
+
+        # The "## Host tools (call with SB.call)" description section: the
+        # shared intro plus one line per tool with its description and
+        # parameter docs. Returns nil when no tools are bound, so the
+        # section is omitted from the description entirely.
+        def tools_section(tools)
+          return nil if tools.empty?
+
+          [
+            "## Host tools (call with SB.call)",
+            HOST_TOOLS_INTRO.rstrip,
+            *tools.values.flat_map { |entry| tool_lines(entry) }
+          ].join("\n")
+        end
+
+        # The "## Server notes" description section: one line per connected
+        # MCP server that sent instructions. Returns nil when no server
+        # did, so the section is omitted from the description entirely.
+        def server_notes_section(mcps)
+          notes = mcps.filter_map do |entry|
+            instructions = entry.instructions.to_s
+            next if instructions.empty?
+
+            "- **#{entry.name}** — #{instructions}"
+          end
+          return nil if notes.empty?
+
+          ["## Server notes", *notes].join("\n")
+        end
+
+        private
+
+        def add_tool(target, bound, name, reserved)
+          instance = normalize_bound_tool(bound)
+          if instance.is_a?(RubyLLM::CodeMode)
+            raise ArgumentError, "cannot bind a RubyLLM::CodeMode inside another CodeMode"
+          end
+
+          name = (name || tool_binding_name(instance)).to_s
+          raise ArgumentError, "tool name must be a non-empty String" if name.empty?
+          if target.key?(name) || reserved.key?(name)
+            raise ArgumentError, "tool name #{name.inspect} is already bound"
+          end
+          total = target.size + reserved.size
+          if total >= MAX_TOOLS
+            raise ArgumentError, "too many tools (#{total + 1}); the limit is #{MAX_TOOLS}"
+          end
+
+          target[name] = ToolEntry.new(name: name, tool: instance)
+        end
+
+        def normalize_bound_tool(bound)
+          if bound.is_a?(Class)
+            unless bound <= RubyLLM::Tool
+              raise ArgumentError, "expected a RubyLLM::Tool (class or instance), got #{bound.inspect}"
+            end
+
+            bound.new
+          elsif bound.is_a?(RubyLLM::Tool)
+            bound
+          else
+            raise ArgumentError, "expected a RubyLLM::Tool (class or instance), got #{bound.inspect}"
+          end
+        end
+
+        # Leaf-only name derivation, like .tool_name — the guest has no use
+        # for the RubyLLM namespace ("MyApp::Tools::Search" binds as "search").
+        def leaf_tool_name(klass)
+          leaf = klass.name.to_s.split('::').last
+          return "" unless leaf
+
+          RubyLLM::Support::Utils.underscore(leaf).delete_suffix('_tool')
+        end
+
+        # Regular tools keep the leaf-only derivation; MCP tools carry their
+        # own (server) name, which is what they bind under inside the sandbox.
+        def tool_binding_name(instance)
+          return instance.name.to_s if instance.respond_to?(:server_name)
+
+          leaf_tool_name(instance.class)
+        end
+
+        def call_bound_tool(entry, args)
+          unless args.is_a?(Hash)
+            raise ArgumentError,
+                  "tool #{entry.name.inspect} expects a hash of arguments, got #{args.class}"
+          end
+
+          normalize_result(entry.tool.call(**args))
+        end
+
+        # MCP tools return RubyLLM::MCP::Result objects; SecurityBox JSON
+        # round-trips handler returns, so give the guest plain data instead
+        # of an inspect string.
+        def normalize_result(value)
+          return value unless defined?(RubyLLM::MCP::Result) && value.is_a?(RubyLLM::MCP::Result)
+
+          { text: value.text, structured: value.structured, error: value.error? }
+        end
+
+        def normalize_mcp_servers(servers, inputs)
+          unless defined?(RubyLLM::MCP)
+            raise ArgumentError,
+                  "RubyLLM::MCP is not available; mcps needs a ruby_llm version with MCP support"
+          end
+
+          servers.map do |server|
+            case server
+            when Class
+              unless server <= RubyLLM::MCP
+                raise ArgumentError, "expected a RubyLLM::MCP (class or instance), got #{server.inspect}"
+              end
+
+              server.new(**inputs)
+            when RubyLLM::MCP
+              server
+            else
+              raise ArgumentError, "expected a RubyLLM::MCP (class or instance), got #{server.inspect}"
+            end
+          end
+        end
+
+        def bind_mcp_server(target_tools, target_mcps, server, reserved)
+          # Fetching .tools connects to the server; snapshot the instructions
+          # right after, so description builds never touch the network.
+          tools_list = server.tools
+          target_mcps << McpEntry.new(
+            name: server.name.to_s, mcp: server, instructions: server.instructions.to_s
+          )
+          tools_list.each { |tool| bind_mcp_tool(target_tools, server, tool, reserved) }
+        end
+
+        def bind_mcp_tool(target_tools, server, tool, reserved)
+          name = tool_binding_name(tool)
+          if (target_tools.key?(name) || reserved.key?(name)) && !server.name.to_s.empty?
+            name = "#{server.name}_#{name}"
+          end
+          add_tool(target_tools, tool, name, reserved)
+        end
+
+        def tool_lines(entry)
+          line = "- `#{entry.name}`"
+          description = entry.tool.description.to_s
+          line += " — #{description}" unless description.empty?
+
+          [line, *param_lines(entry.tool).map { |param| "  #{param}" }]
+        end
+
+        def param_lines(tool)
+          schema = tool.parameters_schema || {}
+          properties = schema["properties"] || {}
+          required = schema["required"] || []
+
+          properties.filter_map do |param, spec|
+            spec ||= {}
+            status = required.include?(param) ? "required" : "optional"
+            line = "- `#{param}` (#{spec["type"] || "string"}, #{status})"
+            description = spec["description"].to_s
+            line += " — #{description}" unless description.empty?
+            line
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/code_mode.rb

@@ -5,7 +5,7 @@ require "security_box"
 
 module RubyLLM
   class CodeMode < Tool
-    VERSION = "0.4.0"
+    VERSION = "0.4.1"
 
     DEFAULT_TIMEOUT_MS = 30_000
     DEFAULT_FUEL_MS = 10_000
@@ -14,36 +14,9 @@ module RubyLLM
     TOOL_USAGE = "tools expects RubyLLM::Tool classes or instances, " \
                  'or "name" => tool pairs'
 
-    FIXED_DESCRIPTION = <<~DESC.freeze
-      Executes Ruby code inside a secure sandbox and returns what it produced.
-
-      ## How to use
-
-      Pass a complete, self-contained Ruby script as the `code` argument. It runs
-      as a fresh Ruby 4.0 process with:
-      - no network, no threads, no subprocesses, and no access to the host
-        filesystem outside the mounted folders listed below;
-      - a private writable scratch directory at `/work`, wiped between executions;
-      - no state carried over between executions, except the contents of
-        read-write folders.
-
-      Capture progress with `puts` (returned as `stdout`) and end with a value or
-      a JSON-serializable data structure (returned as `value`; non-serializable
-      objects come back as their inspect string). If the code raises, the error
-      class, message and backtrace are returned so you can fix and retry.
-    DESC
-
-    HOST_TOOLS_INTRO = <<~DESC.freeze
-      Registered host tools are callable from inside the sandbox. Call one with
-      `SB.call('name', key: value)`: arguments must be JSON-serializable and become
-      the tool's keyword arguments, and the return value is the tool's result
-      (non-serializable values come back as their inspect string). Bad arguments
-      come back as an `error` hash; handler failures raise `SB::ToolError`
-      (rescuable); unknown names raise `SB::UnknownTool`. Limits: at most 1000
-      tool calls per execution and 1 MiB per result.
-    DESC
-
     Mount = Struct.new(:host, :guest, :mode, :description, keyword_init: true) do
+      # The hash shape SecurityBox::Configuration.build expects under
+      # `mounts:`.
       def payload
         { host: host, guest: guest, mode: mode }
       end
@@ -52,450 +25,10 @@ module RubyLLM
     ToolEntry = Struct.new(:name, :tool, keyword_init: true)
 
     McpEntry = Struct.new(:name, :mcp, :instructions, keyword_init: true)
-
-    # Shared registration and rendering core, parameterized by the target
-    # collections: the class-level DSL (tools/mcps) and the instance-level API
-    # (add_tools/add_mcps) run the exact same validation and rollback logic,
-    # only over different hashes/arrays.
-    module Registry
-      EMPTY = {}.freeze
-
-      class << self
-        # Binds tools into `target` (a name => ToolEntry hash). Accepts tool
-        # classes/instances, arrays of them, and "name" => tool pairs (hash
-        # arguments or kwargs). All-or-nothing: on any failure `target` is
-        # restored and the error re-raised. `reserved` lists names already
-        # taken elsewhere (e.g. the class-level bindings for instance
-        # additions); duplicates against it raise.
-        def register_tools(target, args, kwargs, reserved: EMPTY)
-          snapshot = target.dup
-          begin
-            args.flatten(1).each do |item|
-              if item.is_a?(Hash)
-                item.each { |name, bound| add_tool(target, bound, name.to_s, reserved) }
-              else
-                add_tool(target, item, nil, reserved)
-              end
-            end
-            kwargs.each { |name, bound| add_tool(target, bound, name.to_s, reserved) }
-          rescue StandardError
-            target.replace(snapshot)
-            raise
-          end
-          target
-        end
-
-        # Connects MCP server(s) and binds all their tools into
-        # `target_tools`/`target_mcps`. All-or-nothing across both
-        # collections.
-        def register_mcps(target_tools, target_mcps, servers, inputs, reserved: EMPTY)
-          snapshot_tools = target_tools.dup
-          snapshot_mcps = target_mcps.dup
-          flat = servers.flatten(1)
-          raise ArgumentError, "expected at least one RubyLLM::MCP server" if flat.empty?
-          unless inputs.empty? || flat.all?(Class)
-            raise ArgumentError, "inputs only apply when passing RubyLLM::MCP classes"
-          end
-
-          normalize_mcp_servers(flat, inputs).each do |server|
-            bind_mcp_server(target_tools, target_mcps, server, reserved)
-          end
-          target_mcps
-        rescue StandardError
-          target_tools.replace(snapshot_tools)
-          target_mcps.replace(snapshot_mcps)
-          raise
-        end
-
-        def rpc_handlers(tools)
-          tools.transform_values { |entry| ->(args) { call_bound_tool(entry, args) } }
-        end
-
-        def tools_section(tools)
-          return nil if tools.empty?
-
-          [
-            "## Host tools (call with SB.call)",
-            HOST_TOOLS_INTRO.rstrip,
-            *tools.values.flat_map { |entry| tool_lines(entry) }
-          ].join("\n")
-        end
-
-        def server_notes_section(mcps)
-          notes = mcps.filter_map do |entry|
-            instructions = entry.instructions.to_s
-            next if instructions.empty?
-
-            "- **#{entry.name}** — #{instructions}"
-          end
-          return nil if notes.empty?
-
-          ["## Server notes", *notes].join("\n")
-        end
-
-        private
-
-        def add_tool(target, bound, name, reserved)
-          instance = normalize_bound_tool(bound)
-          if instance.is_a?(RubyLLM::CodeMode)
-            raise ArgumentError, "cannot bind a RubyLLM::CodeMode inside another CodeMode"
-          end
-
-          name = (name || tool_binding_name(instance)).to_s
-          raise ArgumentError, "tool name must be a non-empty String" if name.empty?
-          if target.key?(name) || reserved.key?(name)
-            raise ArgumentError, "tool name #{name.inspect} is already bound"
-          end
-          total = target.size + reserved.size
-          if total >= MAX_TOOLS
-            raise ArgumentError, "too many tools (#{total + 1}); the limit is #{MAX_TOOLS}"
-          end
-
-          target[name] = ToolEntry.new(name: name, tool: instance)
-        end
-
-        def normalize_bound_tool(bound)
-          if bound.is_a?(Class)
-            unless bound <= RubyLLM::Tool
-              raise ArgumentError, "expected a RubyLLM::Tool (class or instance), got #{bound.inspect}"
-            end
-
-            bound.new
-          elsif bound.is_a?(RubyLLM::Tool)
-            bound
-          else
-            raise ArgumentError, "expected a RubyLLM::Tool (class or instance), got #{bound.inspect}"
-          end
-        end
-
-        # Leaf-only name derivation, like .tool_name — the guest has no use
-        # for the RubyLLM namespace ("MyApp::Tools::Search" binds as "search").
-        def leaf_tool_name(klass)
-          leaf = klass.name.to_s.split('::').last
-          return "" unless leaf
-
-          RubyLLM::Support::Utils.underscore(leaf).delete_suffix('_tool')
-        end
-
-        # Regular tools keep the leaf-only derivation; MCP tools carry their
-        # own (server) name, which is what they bind under inside the sandbox.
-        def tool_binding_name(instance)
-          return instance.name.to_s if instance.respond_to?(:server_name)
-
-          leaf_tool_name(instance.class)
-        end
-
-        def call_bound_tool(entry, args)
-          unless args.is_a?(Hash)
-            raise ArgumentError,
-                  "tool #{entry.name.inspect} expects a hash of arguments, got #{args.class}"
-          end
-
-          normalize_result(entry.tool.call(**args))
-        end
-
-        # MCP tools return RubyLLM::MCP::Result objects; SecurityBox JSON
-        # round-trips handler returns, so give the guest plain data instead
-        # of an inspect string.
-        def normalize_result(value)
-          return value unless defined?(RubyLLM::MCP::Result) && value.is_a?(RubyLLM::MCP::Result)
-
-          { text: value.text, structured: value.structured, error: value.error? }
-        end
-
-        def normalize_mcp_servers(servers, inputs)
-          unless defined?(RubyLLM::MCP)
-            raise ArgumentError,
-                  "RubyLLM::MCP is not available; mcps needs a ruby_llm version with MCP support"
-          end
-
-          servers.map do |server|
-            case server
-            when Class
-              unless server <= RubyLLM::MCP
-                raise ArgumentError, "expected a RubyLLM::MCP (class or instance), got #{server.inspect}"
-              end
-
-              server.new(**inputs)
-            when RubyLLM::MCP
-              server
-            else
-              raise ArgumentError, "expected a RubyLLM::MCP (class or instance), got #{server.inspect}"
-            end
-          end
-        end
-
-        def bind_mcp_server(target_tools, target_mcps, server, reserved)
-          # Fetching .tools connects to the server; snapshot the instructions
-          # right after, so description builds never touch the network.
-          tools_list = server.tools
-          target_mcps << McpEntry.new(
-            name: server.name.to_s, mcp: server, instructions: server.instructions.to_s
-          )
-          tools_list.each { |tool| bind_mcp_tool(target_tools, server, tool, reserved) }
-        end
-
-        def bind_mcp_tool(target_tools, server, tool, reserved)
-          name = tool_binding_name(tool)
-          if (target_tools.key?(name) || reserved.key?(name)) && !server.name.to_s.empty?
-            name = "#{server.name}_#{name}"
-          end
-          add_tool(target_tools, tool, name, reserved)
-        end
-
-        def tool_lines(entry)
-          line = "- `#{entry.name}`"
-          description = entry.tool.description.to_s
-          line += " — #{description}" unless description.empty?
-
-          [line, *param_lines(entry.tool).map { |param| "  #{param}" }]
-        end
-
-        def param_lines(tool)
-          schema = tool.parameters_schema || {}
-          properties = schema["properties"] || {}
-          required = schema["required"] || []
-
-          properties.filter_map do |param, spec|
-            spec ||= {}
-            status = required.include?(param) ? "required" : "optional"
-            line = "- `#{param}` (#{spec["type"] || "string"}, #{status})"
-            description = spec["description"].to_s
-            line += " — #{description}" unless description.empty?
-            line
-          end
-        end
-      end
-    end
-
-    class << self
-      def tool_name
-        leaf = name.to_s.split('::').last
-        return '' unless leaf
-
-        RubyLLM::Support::Utils.underscore(leaf).delete_suffix('_tool')
-      end
-
-      def mount(source:, dest:, description: nil)
-        add_mount(source, dest, :read_only, description)
-      end
-
-      def mount_rw(source:, dest:, description: nil)
-        add_mount(source, dest, :read_write, description)
-      end
-
-      # Binds RubyLLM::Tool classes/instances as host tools callable from
-      # inside the sandbox via SB.call:
-      #
-      #   tools Weather                     # class (instantiated with .new)
-      #   tools Weather, Echo               # several at once
-      #   tools [Weather, Echo]             # or arrays
-      #   tools "forecast" => Weather       # explicit sandbox name
-      #   tools forecast: Weather           # same, as kwargs
-      #
-      # Without arguments it returns the hash of bound tools. The name is
-      # derived from the class-name leaf (`MyApp::Tools::Weather` binds as
-      # `weather`). Fails fast: every item must be a RubyLLM::Tool (class or
-      # instance), names must be unique and non-empty, and if anything goes
-      # wrong nothing is bound.
-      def tools(*args, **kwargs)
-        return @tools ||= {} if args.empty? && kwargs.empty?
-
-        Registry.register_tools(tools, args, kwargs)
-      end
-
-      # Connects MCP server(s) and binds all their tools into the sandbox:
-      #
-      #   docs = RubyLLM.mcp(url: "https://learn.microsoft.com/api/mcp")
-      #
-      #   class DocsResearch < RubyLLM::CodeMode
-      #     mcps docs                    # instance
-      #     mcps MicrosoftDocs           # or the class (instantiated with .new)
-      #     mcps Linear, user: current_user  # class + declared inputs
-      #     mcps Docs, Github            # several servers / arrays at once
-      #   end
-      #
-      # Without arguments it returns the array of connected servers. Server
-      # tools keep their names (`microsoft_docs_search`, ...); a tool whose
-      # name is already bound is re-bound as "<server>_tool". The servers'
-      # instructions (when the server sends them) are snapshotted at
-      # definition and listed in the description under "## Server notes".
-      # Fails fast: every argument must be a RubyLLM::MCP (class or
-      # instance), and if anything goes wrong — a server that cannot be
-      # reached, a name that stays colliding after prefixing — nothing is
-      # bound.
-      def mcps(*servers, **inputs)
-        return @mcps ||= [] if servers.empty? && inputs.empty?
-
-        Registry.register_mcps(tools, mcps, servers, inputs)
-      end
-
-      def mounts
-        @mounts ||= []
-      end
-
-      def configuration
-        @configuration ||= SecurityBox::Configuration.build(
-          timeout_ms: DEFAULT_TIMEOUT_MS,
-          fuel_ms: DEFAULT_FUEL_MS,
-          mounts: mounts.map(&:payload),
-          rpcs: Registry.rpc_handlers(tools)
-        )
-      end
-
-      # Builds the description the model sees. `extra_tools`/`extra_mcps`
-      # carry the per-instance additions (see #add_tools/#add_mcps); with no
-      # extras the output is identical to the class-level description.
-      def build_description(extra_tools: {}, extra_mcps: [])
-        read_only, read_write = mounts.partition { |m| m.mode == :read_only }
-        [
-          FIXED_DESCRIPTION.rstrip,
-          folder_section("## Read-only folders (readable, never writable)", read_only),
-          folder_section("## Read-write folders (readable and writable)", read_write),
-          Registry.tools_section(tools.merge(extra_tools)),
-          Registry.server_notes_section(mcps + extra_mcps)
-        ].compact.join("\n\n")
-      end
-
-      def description(text = nil)
-        raise ArgumentError,
-              "RubyLLM::CodeMode builds the tool description from the declared mounts; " \
-              "the description is fixed and cannot be set by hand" if text
-
-        build_description
-      end
-
-      def warmup
-        SecurityBox.warmup(timeout_ms: DEFAULT_TIMEOUT_MS, fuel_ms: DEFAULT_FUEL_MS)
-      end
-
-      def inherited(subclass)
-        super
-        subclass.instance_variable_set(:@mounts, mounts.dup)
-        subclass.instance_variable_set(:@tools, tools.dup)
-        subclass.instance_variable_set(:@mcps, mcps.dup)
-        subclass.instance_variable_set(:@configuration, nil)
-      end
-
-      private
-
-      def add_mount(source, dest, mode, description)
-        mounts << Mount.new(
-          host: File.expand_path(source.to_s),
-          guest: dest.to_s,
-          mode: mode,
-          description: description.to_s
-        )
-        SecurityBox::Mounts.normalize(mounts.map(&:payload))
-      rescue SecurityBox::InvalidConfiguration => e
-        mounts.pop
-        raise ArgumentError,
-              "Invalid mount #{source.inspect} => #{dest.inspect}: #{e.message}"
-      end
-
-      def folder_section(header, entries)
-        return nil if entries.empty?
-
-        lines = entries.map do |entry|
-          line = "- `#{entry.guest}`"
-          line += " — #{entry.description}" unless entry.description.empty?
-          line
-        end
-        [header, *lines].join("\n")
-      end
-    end
-
-    parameter :code, description: "Complete Ruby script to execute in the sandbox"
-
-    # Binds more host tools onto this specific instance, on top of the
-    # class-level ones. Accepts the same forms as the class-level .tools:
-    # classes, ready instances, several at once, arrays, and
-    # "name" => tool pairs (hash or kwargs). Instances of the same class are
-    # unaffected, and the additions show up in this instance's description,
-    # so chats built after the call expose them to the model. Fails fast and
-    # rolls everything back; once the first execution happened the sandbox is
-    # already built and further additions raise.
-    def add_tools(*args, **kwargs)
-      ensure_extendable!
-
-      Registry.register_tools(instance_tools, args, kwargs, reserved: self.class.tools)
-    end
-
-    # Connects MCP server(s) onto this specific instance, on top of the
-    # class-level ones. Same forms and fail-fast semantics as .mcps
-    # (instances, classes with their declared inputs, arrays); the servers'
-    # tools bind into this instance only. Raises after the first execution.
-    def add_mcps(*servers, **inputs)
-      ensure_extendable!
-
-      Registry.register_mcps(instance_tools, instance_mcps, servers, inputs, reserved: self.class.tools)
-    end
-
-    def description
-      self.class.build_description(extra_tools: instance_tools, extra_mcps: instance_mcps)
-    end
-
-    def execute(code:)
-      format_result(sandbox.eval(code))
-    rescue SecurityBox::Error => e
-      { status: "sandbox_error", error: { class: e.class.name, message: e.message } }
-    end
-
-    private
-
-    def ensure_extendable!
-      return unless instance_variable_defined?(:@sandbox)
-
-      raise ArgumentError,
-            "cannot add tools or MCP servers after the first execution; " \
-            "create a new #{self.class.name} instance instead"
-    end
-
-    def instance_tools
-      @instance_tools ||= {}
-    end
-
-    def instance_mcps
-      @instance_mcps ||= []
-    end
-
-    # Instances without additions share the class configuration; instances
-    # with additions build their own, merging the class rpcs with the
-    # instance ones (names cannot collide — the registrations validate that).
-    def configuration
-      return self.class.configuration if instance_tools.empty? && instance_mcps.empty?
-
-      @configuration ||= SecurityBox::Configuration.build(
-        timeout_ms: DEFAULT_TIMEOUT_MS,
-        fuel_ms: DEFAULT_FUEL_MS,
-        mounts: self.class.mounts.map(&:payload),
-        rpcs: self.class.configuration.rpcs.merge(Registry.rpc_handlers(instance_tools))
-      )
-    end
-
-    def sandbox
-      @sandbox ||= SecurityBox::Sandbox.new(configuration)
-    end
-
-    def format_result(result)
-      payload = { status: result.status.to_s }
-      case result.status
-      when :ok
-        payload[:value] = result.value
-      when :error
-        payload[:error] = result.error || { "class" => "Unknown", "message" => "guest code failed" }
-      when :timeout
-        payload[:error] = "execution timed out before finishing"
-      when :fuel_exhausted
-        payload[:error] = "execution exceeded its CPU budget before finishing"
-      when :memory_limit
-        payload[:error] = "execution exceeded the sandbox memory limit"
-      when :sandbox_error
-        payload[:error] = "the sandbox failed to run the code"
-      end
-      payload[:stdout] = result.stdout if result.stdout && !result.stdout.empty?
-      payload[:stderr] = result.stderr if result.stderr && !result.stderr.empty?
-      payload
-    end
   end
 end
+
+require_relative "code_mode/description"
+require_relative "code_mode/registry"
+require_relative "code_mode/dsl"
+require_relative "code_mode/instance_methods"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: ruby_llm-code_mode
 version: !ruby/object:Gem::Version
-  version: 0.4.0
+  version: 0.4.1
 platform: ruby
 authors:
 - Juneira
@@ -76,6 +76,10 @@ files:
 - LICENSE
 - README.md
 - lib/ruby_llm/code_mode.rb
+- lib/ruby_llm/code_mode/description.rb
+- lib/ruby_llm/code_mode/dsl.rb
+- lib/ruby_llm/code_mode/instance_methods.rb
+- lib/ruby_llm/code_mode/registry.rb
 licenses:
 - MIT
 metadata: