ruby_llm 2.1.1 → 2.2.0

data/lib/ruby_llm/code_execution/workspace.rb

@@ -0,0 +1,314 @@
+# frozen_string_literal: true
+
+require 'find'
+require 'fileutils'
+require 'tempfile'
+require 'tmpdir'
+
+module RubyLLM
+  class CodeExecution
+    # The directories one run gets: a fresh working directory holding a
+    # link, +attachments+, to read-only copies of the conversation's files
+    # beside it, writable copies of the files named from a chat's FileStore,
+    # links to read-only mounts, separate home and temporary directories, and the
+    # source file. Nothing read-only sits inside the writable working
+    # directory, which Landlock could not keep read-only. Everything is
+    # removed when the run ends. Links to mounts placed in a store's own
+    # folder stay until the last run in that folder ends.
+    class Workspace # :nodoc:
+      # Raised when a run cannot have what it asks for.
+      class Unavailable < StandardError; end
+
+      ATTACHMENTS = 'attachments'
+      IGNORED = %w[__pycache__].freeze
+      OPEN_FLAGS = File::RDONLY | File::BINARY | (defined?(File::NOFOLLOW) ? File::NOFOLLOW : 0)
+      StoreRuns = Struct.new(:active, :placed)
+      STORE_RUNS = Hash.new { |runs, root| runs[root] = StoreRuns.new(0, []) }
+      STORE_RUNS_LOCK = Mutex.new
+      private_constant :StoreRuns, :STORE_RUNS, :STORE_RUNS_LOCK
+
+      attr_reader :directory, :home, :tmp, :inputs, :source, :unmounted
+
+      def self.open(attachments: true)
+        Dir.mktmpdir('ruby_llm-code-') do |root|
+          workspace = new(root, attachments:)
+          yield workspace
+        ensure
+          workspace&.remove_links
+          workspace&.close_outputs
+          unlock(root)
+        end
+      end
+
+      def self.unlock(root)
+        Find.find(root) do |path|
+          File.chmod(0o700, path) if File.lstat(path).directory?
+        rescue SystemCallError
+          next
+        end
+      end
+
+      def self.filename(name)
+        path = name.to_s.tr('\\', '/')
+        return path unless path.empty? || path.split('/', -1).intersect?(['', '.', '..'])
+
+        base = File.basename(path)
+        ['', '.', '..'].include?(base) ? 'file' : base
+      end
+
+      def self.invalid_name(name)
+        "#{name.scrub} (name is not valid UTF-8)"
+      end
+
+      # Returns each attachment's name under +attachments/+, newest first:
+      # the newest keeps its name, as it would on a disk, and an older one
+      # stays readable beside it. Names compare without case, for the
+      # filesystems that do, and files with equal content share one.
+      def self.input_names(attachments)
+        attachments = attachments.reject(&:provider_file?)
+        folders = input_folders(attachments)
+        taken = {}
+        attachments.reverse.each_with_object({}.compare_by_identity) do |attachment, names|
+          name = filename(attachment.filename)
+          choice = numbered_names(name).find do |candidate|
+            (candidate == name || !folders.include?(candidate.downcase)) &&
+              !different?(taken[candidate.downcase], attachment)
+          end
+          taken[choice.downcase] ||= [choice, attachment]
+          names[attachment] = taken[choice.downcase].first
+        end
+      end
+
+      def self.different?(holder, attachment)
+        holder && !holder.last.equal?(attachment) && holder.last.content.b != attachment.content.b
+      end
+      private_class_method :different?
+
+      def self.input_folders(attachments)
+        attachments.flat_map do |attachment|
+          parts = filename(attachment.filename).downcase.split('/')
+          (1...parts.size).map { |size| parts.first(size).join('/') }
+        end
+      end
+      private_class_method :input_folders
+
+      def self.numbered_names(name)
+        extension = File.extname(name)
+        stem = name.delete_suffix(extension)
+        [name].chain((2..).lazy.map { |number| "#{stem}-#{number}#{extension}" })
+      end
+      private_class_method :numbered_names
+
+      def initialize(root, attachments: true)
+        @directory = File.join(root, 'work')
+        @inputs = File.join(root, 'inputs')
+        @home = File.join(root, 'home')
+        @tmp = File.join(root, 'tmp')
+        @source = File.join(root, 'source')
+        [@directory, @inputs, @home, @tmp, @source].each { |path| FileUtils.mkdir_p(path) }
+        @attachments = File.join(@directory, ATTACHMENTS)
+        File.symlink(@inputs, @attachments) if attachments
+        @unmounted = []
+        @opened = []
+      end
+
+      def self.mount_path(name)
+        parts = name.to_s.tr('\\', '/').split('/')
+        if parts.empty? || name.to_s.start_with?('/') || parts.intersect?(['', '.', '..']) || parts.first == ATTACHMENTS
+          raise ArgumentError, "#{name.inspect} is not a path inside the working directory"
+        end
+
+        parts.join('/')
+      end
+
+      # Links each mount into +into+ and returns the targets. Into a folder
+      # other than the working directory, a link reuses one already there
+      # to the same target, leaves a path that holds anything else alone
+      # (see #unmounted), and is removed with the folders it made by
+      # #remove_links.
+      def mount(mounts, into: directory)
+        return place_all(mounts, into) if into == directory
+
+        STORE_RUNS_LOCK.synchronize do
+          @store_root = into
+          STORE_RUNS[into].active += 1
+          place_all(mounts, into)
+        end
+      end
+
+      # Ends this run in a store's folder. The last run there removes what
+      # the runs placed.
+      def remove_links
+        return unless @store_root
+
+        STORE_RUNS_LOCK.synchronize do
+          runs = STORE_RUNS[@store_root]
+          runs.active -= 1
+          next unless runs.active.zero?
+
+          STORE_RUNS.delete(@store_root)
+          runs.placed.reverse_each { |path, target| remove_placed(path, target) }
+        end
+      end
+
+      def write_source(language, code)
+        File.join(source, language.filename).tap { |path| File.write(path, code) }
+      end
+
+      def link_packages(node_modules)
+        File.symlink(node_modules, File.join(source, 'node_modules'))
+      end
+
+      def copy_in(attachments)
+        self.class.input_names(attachments).to_a.reverse_each do |attachment, name|
+          path = File.join(inputs, name)
+          make_room(path, inputs)
+          File.binwrite(path, attachment.content.b)
+        end
+        {}
+      end
+
+      # Copies +files+, StoredFiles of the conversation's files, to
+      # +attachments/+ at their paths in +store+, where the model finds them.
+      def copy_stored_inputs(store, files)
+        files.each do |file|
+          path = File.join(inputs, file.path)
+          make_room(path, inputs)
+          File.binwrite(path, store.read(file.path))
+        end
+      end
+
+      def copy_files(store, files)
+        files.each_with_object({}) do |file, copied|
+          next if file.path.split('/').any? { |part| part.start_with?('.') }
+
+          content = store.read(file.path)
+          path = File.join(directory, file.path)
+          make_room(path, directory)
+          File.binwrite(path, content)
+          copied[file.path] = [content.bytesize, FileStore.digest(content)]
+        end
+      end
+
+      def outputs(max_files:, max_size:, unchanged: {})
+        files = []
+        skipped = []
+        output_paths.each do |path, size|
+          name = path.delete_prefix(File.join(directory, ''))
+          next skipped << self.class.invalid_name(name) unless name.valid_encoding?
+          next if unchanged?(path, size, unchanged[name])
+          next skipped << "#{name} (larger than #{max_size} bytes)" if size > max_size
+          next skipped << "#{name} (more than #{max_files} files)" if files.size >= max_files
+
+          files << Attachment.new(open_output(path), filename: name)
+        rescue SystemCallError
+          skipped << "#{name} (unreadable)"
+        end
+        [files, skipped]
+      end
+
+      # Returns +files+, from #outputs, copied to tempfiles that outlive
+      # the workspace, one file at a time.
+      def keep(files)
+        files.map do |file|
+          copy = Tempfile.new(['ruby_llm-output-', File.extname(file.filename)], binmode: true)
+          IO.copy_stream(file.source, copy, nil, 0)
+          Attachment.new(copy, filename: file.filename, mime_type: file.mime_type)
+        end
+      end
+
+      def close_outputs
+        @opened.each(&:close)
+      end
+
+      private
+
+      def place_all(mounts, into)
+        mounts.filter_map do |name, path|
+          target = File.expand_path(path)
+          link = File.join(into, self.class.mount_path(name))
+          next target if into == directory ? place(link, target) : place_in_store(link, target, into)
+
+          @unmounted << self.class.mount_path(name)
+          nil
+        end
+      end
+
+      def place(link, target)
+        FileUtils.mkdir_p(File.dirname(link))
+        File.symlink(target, link)
+      end
+
+      def place_in_store(link, target, root)
+        make_parents(File.dirname(link), root)
+        return true if File.symlink?(link) && File.readlink(link) == target
+
+        File.symlink(target, link)
+        STORE_RUNS[root].placed << [link, target]
+        true
+      rescue SystemCallError
+        false
+      end
+
+      def make_parents(folder, root)
+        return if folder == root
+
+        make_parents(File.dirname(folder), root)
+        return if File.directory?(folder) && !File.symlink?(folder)
+
+        Dir.mkdir(folder)
+        STORE_RUNS[root].placed << [folder, nil]
+      end
+
+      def remove_placed(path, target)
+        if target
+          File.delete(path) if File.symlink?(path) && File.readlink(path) == target
+        else
+          Dir.rmdir(path)
+        end
+      rescue SystemCallError
+        nil
+      end
+
+      def make_room(path, base)
+        parent = base
+        File.dirname(path).delete_prefix(base).split('/').reject(&:empty?).each do |part|
+          parent = File.join(parent, part)
+          FileUtils.rm_f(parent) unless File.directory?(parent)
+        end
+        FileUtils.mkdir_p(File.dirname(path))
+        FileUtils.rm_rf(path) if File.directory?(path)
+      end
+
+      def output_paths
+        paths = []
+        Find.find(directory) do |path|
+          name = File.basename(path)
+          Find.prune if path == @attachments || IGNORED.include?(name) || (path != directory && name.start_with?('.'))
+          size = output_size(path)
+          paths << [path, size] if size
+        end
+        paths.sort
+      end
+
+      def output_size(path)
+        stat = File.lstat(path)
+        return stat.size if stat.file?
+
+        0 if stat.directory? && !(stat.readable? && stat.executable?)
+      rescue SystemCallError
+        0
+      end
+
+      def unchanged?(path, size, copied)
+        return false unless copied&.first == size
+
+        File.open(path, OPEN_FLAGS) { |io| copied.last == FileStore.digest(io) }
+      end
+
+      def open_output(path)
+        File.new(path, OPEN_FLAGS).tap { |io| @opened << io }
+      end
+    end
+  end
+end