omakase-agents 0.2.0 → 0.4.0

data/lib/omakase/agent.rb

@@ -4,6 +4,11 @@ module Omakase
   # Fields are state, methods are what the model can call, `generates` declares
   # the methods the model implements.
   class Agent
+    # The generations this thread is inside, so one cannot re-enter itself.
+    RUNNING = :omakase_running
+    # Each level opens a chat with a fresh tool budget, so depth is what bounds the cost.
+    # ponytail: a constant; a setting when a real tree needs more.
+    MAX_DEPTH = 10
     class << self
       # The model and any RubyLLM chat option. Naming a provider takes the model
       # id on trust, since providers like OpenRouter or Ollama serve ids that are
@@ -29,9 +34,14 @@ module Omakase
 
       # An MCP server's tools, as methods on the agent. Options are passed to
       # `ruby_llm-mcp` verbatim: `mcp :files, transport_type: :stdio, config: {command: "npx", …}`.
-      def mcp(name, **options)
-        require "ruby_llm/mcp"
-        MCP.attach(self, RubyLLM::MCP.add_client(name: name.to_s, **options))
+      # The server opens on the first instance, not at class load.
+      def mcp(name, **options) = MCP.defer(self, name, options)
+
+      # Here rather than in `initialize`, which a subclass may override and
+      # never call super from — and then generate with none of its tools.
+      def new(...)
+        MCP.ensure(self)
+        super
       end
 
       # A skill directory — a SKILL.md with YAML front matter. Its description
@@ -53,8 +63,9 @@ module Omakase
       end
 
       # Without a prompt, the method name is the prompt. A block instead of a
-      # string is a prompt read at call time, on the agent.
-      def generates(name, prompt = nil, returns: nil, strategy: nil, model: nil, &schema)
+      # string is a prompt read at call time, on the agent. `takes:` names the
+      # keyword arguments, and then Ruby checks them.
+      def generates(name, prompt = nil, takes: nil, returns: nil, strategy: nil, model: nil, &schema)
         # Redeclaring an inherited generation is how a subclass specialises one.
         # Landing on a method you wrote is not that, and would replace it unseen.
         if Capabilities.names(self).include?(name) && !generations.key?(name)
@@ -72,7 +83,7 @@ module Omakase
           strategy: Strategies.fetch(strategy || self.strategy),
           model:
         )
-        define_method(name) { |**inputs| generate(name, inputs) }
+        define_generation_method(name, takes)
         define_singleton_method(name) { |**inputs| new.public_send(name, **inputs) }
       end
 
@@ -84,6 +95,25 @@ module Omakase
 
       private
 
+      # Named inputs become a real signature, so a missing or misspelled argument
+      # is an ArgumentError at the call rather than noise in a prompt — and the
+      # model reads the names too, instead of `**inputs`.
+      def define_generation_method(name, takes)
+        return define_method(name) { |**inputs| generate(name, inputs) } if takes.nil?
+
+        keywords = Array(takes)
+        bad = [name.to_s.chomp("?").chomp("!"), *keywords].reject { |word| /\A[a-z_]\w*\z/.match?(word.to_s) }
+        raise Error, "#{self}##{name}: takes: needs plain keyword names, got #{bad.inspect}" if bad.any?
+
+        class_eval <<~RUBY, __FILE__, __LINE__ + 1
+          def #{name}(#{keywords.map { |key| "#{key}:" }.join(", ")}, with: nil)
+            inputs = {#{keywords.map { |key| "#{key}: #{key}" }.join(", ")}}
+            inputs[:with] = with unless with.nil?
+            generate(:#{name}, inputs)
+          end
+        RUBY
+      end
+
       def humanize(name)
         text = name.to_s.tr("_", " ").capitalize
         text.end_with?("?", "!") ? text : "#{text}."
@@ -102,6 +132,7 @@ module Omakase
         subclass.strategy(@strategy) if @strategy
         subclass.generations.merge!(generations)
         subclass.descriptions.merge!(descriptions)
+        Skills.attach_core(subclass)
       end
     end
 
@@ -113,7 +144,7 @@ module Omakase
     # A fresh conversation per call — two threads calling one agent must not
     # share a mutable chat. What carries between calls is the object's own state.
     # Overrides land on top of the class's options; an injected chat ignores them.
-    def chat(**overrides) = @chat || RubyLLM.chat(**self.class.chat_options.merge(overrides))
+    def chat(**overrides) = @chat || Omakase.chat_factory.call(**self.class.chat_options.merge(overrides))
 
     # That state, as the model should read it: rebuilt on every call, and added
     # to the class's instructions. Override it to remember anything.
@@ -150,12 +181,48 @@ module Omakase
 
     def generate(name, inputs)
       generation = self.class.generations.fetch(name)
-      Omakase.emit(:generation, agent: self, name:, inputs:)
-      value = generation.strategy.call(Request.new(agent: self, generation:, inputs:))
-      Omakase.emit(:answer, agent: self, name:, value:)
-      value
-    rescue RubyLLM::Error, RubyLLM::ConfigurationError, RubyLLM::ModelNotFoundError => e
-      raise ProviderError, "#{self.class}##{name}: #{e.message}"
+      running = (Thread.current[RUNNING] ||= [])
+      key = [object_id, name]
+
+      # Generated code can see this method and call it. Each nested call opens its
+      # own chat with its own tool budget, so the budget would bound nothing.
+      raise Error, "#{self.class}##{name} is already running — it cannot call itself" if running.any? { |frame| frame[:key] == key }
+      raise Error, "#{self.class}##{name}: generations nest deeper than #{MAX_DEPTH}" if running.size >= MAX_DEPTH
+
+      frame = {key:, inputs:}
+      running.push(frame)
+      begin
+        Omakase.emit(:generation, agent: self, name:, inputs:)
+        value = generation.strategy.call(Request.new(agent: self, generation:, inputs:))
+        Omakase.emit(:answer, agent: self, name:, value:)
+        value
+      rescue RubyLLM::Error, RubyLLM::ConfigurationError, RubyLLM::ModelNotFoundError => e
+        error = ProviderError.new("#{self.class}##{name}: #{e.message}")
+        Omakase.emit(:error, agent: self, name:, error:)
+        raise error
+      rescue => e
+        Omakase.emit(:error, agent: self, name:, error: e)
+        raise
+      ensure
+        running.delete(frame)
+      end
     end
+
+    # Where generated code runs during a generation: self is the agent and the
+    # inputs are locals, which then last for the rest of that generation.
+    def omakase_scope
+      frame = Thread.current[RUNNING]&.reverse_each&.find { |entry| entry[:key].first == object_id }
+      return unless frame
+
+      frame[:scope] ||= omakase_blank_binding.tap do |scope|
+        frame[:inputs].except(:with).each do |name, value|
+          scope.local_variable_set(name, value) if /\A[a-z_]\w*\z/.match?(name.to_s)
+        end
+      end
+    end
+
+    # A string instance_eval, so a `def` in generated code lands on this object
+    # alone — and no local of the caller leaks in.
+    def omakase_blank_binding = instance_eval("binding", __FILE__, __LINE__)
   end
 end

data/lib/omakase/capabilities.rb

@@ -18,10 +18,15 @@ module Omakase
     end
 
     def entry(agent_class, name)
-      signature = "#{name}(#{parameters(agent_class.instance_method(name))})"
+      method = agent_class.instance_method(name)
+      signature = "#{name}(#{parameters(method)})"
       # A prompt written as a block needs an instance to read; `describe` it instead.
       prompt = agent_class.generations[name]&.prompt
-      description = agent_class.descriptions[name] || (prompt unless prompt.is_a?(Proc))
+      # Look up on the method's owner so a late attach on a parent still
+      # documents the tool for subclasses created before that generate.
+      owner = method.owner
+      description = (owner.descriptions[name] if owner.respond_to?(:descriptions)) ||
+        (prompt unless prompt.is_a?(Proc))
       description ? "#{signature} — #{description}" : signature
     end
 

data/lib/omakase/doc.rb

@@ -13,7 +13,23 @@ module Omakase
     # and the methods you wrote.
     def boundary = defined?(ActiveRecord::Base) ? [*CORE, ActiveRecord::Base] : CORE
 
-    def of(object) = [object.class.to_s, *signatures(object), *state(object)].join("\n")
+    def of(object)
+      return of_class(object) if object.is_a?(Module)
+
+      [object.class.to_s, *signatures(object.class) { |name| object.method(name) }, *state(object)].join("\n")
+    end
+
+    # A class, not an instance: what one would have. The model asks this before it
+    # builds an object of a type it has only been told the name of.
+    def of_class(klass)
+      [klass.to_s, *signatures(klass) { |name| klass.instance_method(name) }, *columns(klass)].join("\n")
+    end
+
+    def columns(klass)
+      return [] unless klass.respond_to?(:column_names)
+
+      klass.column_names.map { |name| "  #{name}" }
+    end
 
     # An object that answers `attributes` says what it holds better than its
     # instance variables do — and a record's columns are state, not API.
@@ -24,14 +40,14 @@ module Omakase
 
     def ivars(object) = object.instance_variables.to_h { |name| [name, object.instance_variable_get(name)] }
 
-    def signatures(object)
-      object.class.ancestors
+    def signatures(klass, &getter)
+      klass.ancestors
         .take_while { |mod| !boundary.include?(mod) }
         .reject { |mod| mod.to_s.end_with?("GeneratedAttributeMethods") }
         .flat_map { |mod| mod.public_instance_methods(false) }
         .uniq.sort
         .reject { |name| name.match?(/\A_|_associated_records_for_/) }
-        .map { |name| "  #{name}(#{Capabilities.parameters(object.method(name))})" }
+        .map { |name| "  #{name}(#{Capabilities.parameters(getter.call(name))})" }
     end
   end
 end

data/lib/omakase/executor.rb

@@ -2,7 +2,7 @@
 
 module Omakase
   # Runs model-written Ruby in the agent's own context.
-  # ponytail: instance_eval is not a sandbox — see Omakase.executor to swap it.
+  # ponytail: instance_eval is not a sandbox — Subprocess isolates a crash, not File.
   module Executor
     SOURCE = "(generated)"
     RESULT = :omakase_result
@@ -11,33 +11,52 @@ module Omakase
     TRACE = /\A#{Regexp.escape(SOURCE)}:\d+/
     MAX_OUTPUT = 4_000
 
-    # What `finish(value)` handed back: the answer as a Ruby value, not as text.
-    Answer = Data.define(:value)
+    # What `finish(value)` handed back: the answer as a Ruby value, not as text,
+    # and whatever the code printed on the way there. `printed` defaults, so a
+    # replacement executor that only knows the value still satisfies the seam.
+    Answer = Data.define(:value, :printed) do
+      def initialize(value:, printed: "") = super
+    end
 
     module_function
 
     def call(agent, code, timeout: TIMEOUT)
       printed = StringIO.new
       answer = catch(RESULT) do
-        value = capturing(printed) { Timeout.timeout(timeout) { agent.instance_eval(code, SOURCE, 1) } }
+        value = capturing(printed) { Timeout.timeout(timeout) { evaluate(agent, code) } }
         return observation([printed.string.chomp, "=> #{value.inspect}"])
       end
-      Answer.new(value: answer)
+      Answer.new(value: answer, printed: printed.string.chomp)
     rescue ScriptError, StandardError => e
       observation([printed.string.chomp, failure(e, code)])
+    rescue SystemExit
+      # In process, exit would end the host. exit! cannot be caught: Subprocess covers that.
+      observation([printed.string.chomp, "exit is not allowed — answer with finish(value)"])
+    end
+
+    # Inside a generation the inputs are locals; outside one, plain instance_eval.
+    def evaluate(agent, code)
+      scope = agent.respond_to?(:omakase_scope, true) && agent.send(:omakase_scope)
+      scope ? scope.eval(code, SOURCE, 1) : agent.instance_eval(code, SOURCE, 1)
     end
 
-    # The model can only fix what it can locate, so point at the line.
+    # The model can only fix what it can locate, so point at the line. A
+    # SyntaxError has no generated frame to point from.
     def failure(error, code)
-      line = error.backtrace&.grep(TRACE)&.first&.slice(/:(\d+)/, 1)&.to_i
-      source = code.lines[line - 1]&.strip if line&.positive?
-      ["#{error.class}: #{error.message}", ("line #{line}: #{source}" if source)].compact.join("\n")
+      message = "#{error.class}: #{error.message}"
+      line = error.backtrace.grep(TRACE).first&.slice(/:(\d+)/, 1)&.to_i
+      return message unless line
+
+      "#{message}\nline #{line}: #{code.lines[line - 1].to_s.strip}"
     end
 
+    # Head and tail: the error or the value comes last, and it is what the model needs next.
     def observation(parts)
       text = parts.reject(&:empty?).join("\n")
-      text = "#{text[0, MAX_OUTPUT]}\n… (truncated)" if text.length > MAX_OUTPUT
-      text.empty? ? "(no output)" : text
+      return text if text.length <= MAX_OUTPUT
+
+      half = MAX_OUTPUT / 2
+      "#{text[0, half]}\n… (#{text.length - MAX_OUTPUT} characters truncated)\n#{text[-half..]}"
     end
 
     # Thread-local, so concurrent agents never share a buffer. Agent#puts reads it.
@@ -48,5 +67,127 @@ module Omakase
     ensure
       Thread.current[OUTPUT] = previous
     end
+
+    # Generated code runs in a child process so a timeout, a crash, or a
+    # runaway loop cannot take the parent with it. The child is a copy of
+    # this process — it can still reach ActiveRecord, ENV, and the disk.
+    # That is isolation of fate, not of capability. Untrusted input still
+    # belongs to :predict.
+    #
+    # Ivars written in the child are marshalled back one at a time, so a
+    # generation's second tool call sees what the first one set. Methods
+    # the model defined on the object die with the child.
+    module Subprocess
+      module_function
+
+      def call(agent, code, timeout: TIMEOUT)
+        IO.pipe(binmode: true) do |reader, writer|
+          pid = fork do
+            reader.close
+            # Own process group, so a timeout can kill grandchildren too.
+            Process.setsid
+            # Parent owns the deadline; Timeout here would race it.
+            payload = pack(agent, Executor.call(agent, code, timeout: nil))
+            writer.write([payload.bytesize].pack("N"), payload)
+          ensure
+            exit! 0
+          end
+          writer.close
+          collect(reader, agent, pid, clock + timeout)
+        end
+      end
+
+      def collect(reader, agent, pid, deadline)
+        payload = read_packet(reader, deadline)
+        stop(pid) if payload == :timeout
+        status = reap(pid)
+        case payload
+        when :timeout then "execution timed out"
+        when :eof then "child process #{fate(status)}"
+        else unpack(agent, payload)
+        end
+      end
+
+      def pack(agent, result)
+        kept, dropped = agent.marshal_dump.partition { |_, value| marshalable?(value) }
+        result = note_dropped(result, dropped.map(&:first))
+        Marshal.dump({result: carry(result), state: kept.to_h})
+      end
+
+      def marshalable?(value)
+        Marshal.dump(value)
+        true
+      rescue TypeError
+        false
+      end
+
+      # A dropped ivar turns the answer into an observation: silent state loss
+      # would leave the next tool call reasoning about a value that is gone.
+      def note_dropped(result, dropped)
+        return result if dropped.empty?
+
+        prior = result.is_a?(Answer) ? [result.printed, "finish #{result.value.inspect}"].reject(&:empty?).join("\n") : result
+        Executor.observation([prior, "cannot keep #{dropped.join(", ")} across the process boundary"])
+      end
+
+      # Only an Answer can fail here — an observation is a String.
+      def carry(result)
+        return result if marshalable?(result)
+
+        "cannot return #{result.value.class} across the process boundary"
+      end
+
+      def unpack(agent, payload)
+        packet = Marshal.load(payload)
+        agent.marshal_load(packet[:state])
+        packet[:result]
+      rescue ArgumentError, TypeError => e
+        "#{e.message}: a class defined in generated code does not survive the process boundary"
+      end
+
+      # Length-prefixed, so a leftover write-end cannot hang the parent.
+      def read_packet(io, deadline)
+        header = read_exactly(io, 4, deadline)
+        return header if header.is_a?(Symbol)
+
+        read_exactly(io, header.unpack1("N"), deadline)
+      end
+
+      # select is exact for a pipe, so readpartial cannot block past the deadline.
+      def read_exactly(io, n, deadline)
+        buf = "".b
+        while buf.bytesize < n
+          return :timeout unless IO.select([io], nil, nil, [deadline - clock, 0].max)
+
+          buf << io.readpartial(n - buf.bytesize)
+        end
+        buf
+      rescue EOFError
+        :eof
+      end
+
+      def clock = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+
+      def fate(status)
+        return "was killed" if status.nil? || status.signaled?
+        return "ended without an answer" if status.success?
+
+        "exited #{status.exitstatus}"
+      end
+
+      # The child led its own group unless the deadline beat it to setsid.
+      def stop(pid)
+        Process.kill("KILL", (Process.getpgid(pid) == pid) ? -pid : pid)
+      rescue Errno::ESRCH
+        nil
+      end
+
+      # ECHILD: the host reaps children itself, with a CHLD trap.
+      def reap(pid)
+        Process.wait2(pid).last
+      rescue Errno::ECHILD
+        nil
+      end
+    end
   end
 end

data/lib/omakase/fake_chat.rb

@@ -9,9 +9,22 @@ module Omakase
   # The script receives the chat, so it can drive the tool the way a model would:
   #
   #   Omakase::FakeChat.new { |chat| chat.run("finish(42)") }
+  #
+  # Or one reply per model turn, in order; a turn past the last one raises:
+  #
+  #   Omakase::FakeChat.replies("prose, not JSON", {"result" => 42})
   class FakeChat
     Response = Struct.new(:content)
 
+    def self.replies(*replies)
+      new do |chat|
+        raise Error, "FakeChat: no scripted reply left for turn #{chat.tasks.size}" if replies.empty?
+
+        reply = replies.shift
+        reply.respond_to?(:call) ? reply.call(chat) : reply
+      end
+    end
+
     attr_reader :instructions, :schema, :tools, :tasks, :attachments
 
     def initialize(&script)
@@ -20,20 +33,32 @@ module Omakase
       @tools = []
       @tasks = []
       @attachments = []
+      @complete = true
     end
 
-    def with_instructions(text) = tap { @instructions << text }
+    # A log of every call — the real chat replaces unless `append:`.
+    def with_instructions(text, **) = tap { @instructions << text }
 
     def with_schema(schema) = tap { @schema = schema }
 
-    def with_tool(tool, **) = tap { @tools << tool }
+    def with_tools(*tools) = tap { @tools.concat(tools) }
 
-    def ask(task, with: nil)
+    def ask(task, with: nil) = ask_later(task, with:).step
+
+    def ask_later(task, with: nil)
       @tasks << task
       @attachments << with if with
+      tap { @complete = false }
+    end
+
+    # The whole script is one step: it answers, so the chat is then complete.
+    def step
+      @complete = true
       Response.new(@script.call(self))
     end
 
+    def complete? = @complete
+
     # Run code the way the model would, through the agent's one tool.
     def run(code) = tools.fetch(0).call(code:)
   end

data/lib/omakase/mcp.rb

@@ -4,19 +4,91 @@ module Omakase
   # An MCP server's tools, as methods on the agent — so generated code calls a
   # remote tool the same way it calls anything else the agent exposes.
   module MCP
+    LOCKS = Mutex.new
+    # A tool name landing on a method is your bug, not a sidecar being down —
+    # and both arrive as an Error out of the same call.
+    Clash = Class.new(Error)
+
+    class << self
+      def client_factory=(factory)
+        unless factory.nil? || factory.respond_to?(:call)
+          raise Error, "client_factory must answer call, got #{factory.class}"
+        end
+
+        @client_factory = factory
+      end
+
+      def client_factory
+        @client_factory || ->(name, **options) {
+          require "ruby_llm/mcp"
+          RubyLLM::MCP.add_client(name: name.to_s, **options)
+        }
+      end
+    end
+
     module_function
 
+    def defer(agent_class, name, options)
+      servers(agent_class)[name.to_sym] = options
+    end
+
+    def ensure(agent_class)
+      return unless pending?(agent_class)
+
+      agent_class.ancestors.take_while { |mod| mod != Agent }.reverse_each do |klass|
+        next unless klass.is_a?(Class)
+
+        mutex_for(klass).synchronize { attach_pending(klass) }
+      end
+    end
+
+    def pending?(agent_class)
+      agent_class.ancestors.take_while { |mod| mod != Agent }.any? do |mod|
+        mod.is_a?(Class) && (servers(mod).keys - attached(mod)).any?
+      end
+    end
+
+    def attach_pending(klass)
+      servers(klass).each do |name, options|
+        next if attached(klass).include?(name)
+
+        # A down sidecar stays unattached — connecting and listing its tools are
+        # the same trip. The generate still runs. Next ensure retries. A name
+        # clash is not the sidecar being down, it is your bug, so it raises.
+        begin
+          attach(klass, client_factory.call(name, **options))
+        rescue Clash
+          raise
+        rescue => e
+          Omakase.emit(:mcp, agent: klass, name:, error: e)
+          next
+        end
+        attached(klass) << name
+      end
+    end
+
     def attach(agent_class, client)
+      defined = []
       client.tools.each do |tool|
         name = method_name(tool)
         # A remote tool list must not quietly shadow a capability the agent already has.
-        raise Error, "#{agent_class} already has ##{name}" if Capabilities.names(agent_class).include?(name)
+        raise Clash, "#{agent_class} already has ##{name}" if Capabilities.names(agent_class).include?(name)
 
         agent_class.describe(description(tool))
         # nil is how a model leaves an argument out; MCP servers reject it.
         agent_class.define_method(name) { |**arguments| MCP.result(tool.execute(**arguments.compact)) }
+        defined << name
       end
       client
+    rescue
+      # Else the next ensure dies on "already has #name".
+      defined.each do |name|
+        agent_class.send(:remove_method, name)
+        agent_class.descriptions.delete(name)
+      end
+      agent_class.instance_variable_set(:@pending_description, nil)
+      client.close if client.respond_to?(:close)
+      raise
     end
 
     # Tool names may hold characters a Ruby method name cannot.
@@ -39,5 +111,22 @@ module Omakase
       text = tool.description.to_s.gsub(/\s+/, " ").strip
       [text, ("Arguments — #{arguments.join(", ")}" if arguments.any?)].compact.join(" ")
     end
+
+    def mutex_for(klass)
+      mutex = klass.instance_variable_get(:@mcp_mutex)
+      return mutex if mutex
+
+      LOCKS.synchronize do
+        klass.instance_variable_get(:@mcp_mutex) || klass.instance_variable_set(:@mcp_mutex, Mutex.new)
+      end
+    end
+
+    def servers(klass)
+      klass.instance_variable_get(:@mcp_servers) || klass.instance_variable_set(:@mcp_servers, {})
+    end
+
+    def attached(klass)
+      klass.instance_variable_get(:@mcp_attached) || klass.instance_variable_set(:@mcp_attached, [])
+    end
   end
 end

data/lib/omakase/request.rb

@@ -7,17 +7,21 @@ module Omakase
 
     def schema = generation.schema
 
-    def instructions = [agent.class.instructions, agent.context].reject { |text| text.to_s.empty? }.join("\n\n")
+    def instructions = agent.class.instructions
+
+    # Rebuilt on every call, so it goes after what a provider can cache.
+    def context = agent.context.to_s
 
     # `with:` is reserved: files for the model to look at, passed through to
     # RubyLLM's `ask(with:)` as attachments rather than rendered into the text.
     def attachments = inputs[:with]
 
-    def task
+    # With `preview:`, the model's code holds the inputs, so the prompt only shows them.
+    def task(preview: false)
       arguments = inputs.except(:with)
       return prompt if arguments.empty?
 
-      lines = arguments.map { |name, value| "- #{name}: #{value.inspect}" }
+      lines = arguments.map { |name, value| "- #{name}: #{preview ? shorten(value.inspect) : value.inspect}" }
       "#{prompt}\n\nInputs:\n#{lines.join("\n")}"
     end
 
@@ -27,5 +31,13 @@ module Omakase
       text = generation.prompt
       text.is_a?(Proc) ? agent.instance_exec(&text) : text
     end
+
+    # Past the limit, code_act shows the start of an input: the whole value is a local.
+    # ponytail: inspects the whole value, then cuts; a bounded printer if a huge input shows up in a profile.
+    def shorten(text, limit = 500)
+      return text if text.length <= limit
+
+      "#{text[0, limit]}… (#{text.length} characters — the whole value is in the local)"
+    end
   end
 end