ruby_llm-ag_ui 0.1.0

data/lib/ruby_llm/ag_ui/sse_writer.rb

@@ -0,0 +1,119 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module AgUI
+    # Writes AG-UI events to any IO-like stream as Server-Sent Events.
+    #
+    # The stream is whatever the host is writing a response to: a Rack response
+    # body, an ActionController::Live response, a socket, a StringIO in a test.
+    # The gem does not own the transport, so it does not open connections,
+    # buffer frames or choose how a host relays a background job to a client.
+    # What it does own is the ordering and the framing of the events it writes.
+    #
+    # A dead stream is reported as a StreamError and the writer stays closed.
+    # The reason is a transport truncation, not a run outcome: it is never
+    # evidence that the work finished, succeeded or was cancelled.
+    class SSEWriter
+      # The media type an SSE response must be served as.
+      CONTENT_TYPE = 'text/event-stream'
+
+      # SSE responses must not be buffered by an intermediary.
+      CACHE_CONTROL = 'no-cache'
+
+      # Keep-alive for the proxies that need it. Puma itself does not.
+      CONNECTION = 'keep-alive'
+
+      # A response that streams must not be compressed, because compression
+      # buffers and withholds the frames the client is waiting for.
+      NO_TRANSFORM = 'no'
+
+      # The number of events successfully written.
+      attr_reader :events_written
+
+      # The bytes successfully written.
+      attr_reader :bytes_written
+
+      def initialize(stream, encoder: nil, logger: nil)
+        @stream = stream
+        # The SDK encoder always writes SSE; its accept parameter only names the
+        # media type, and this gem always serves text/event-stream.
+        @encoder = encoder || ::AgUiProtocol::EventEncoder.new
+        @logger = logger || RubyLLM::AgUI.config.logger
+        @events_written = 0
+        @bytes_written = 0
+        @closed = false
+      end
+
+      # Whether an Accept header means the client can read an SSE stream.
+      #
+      # A host calls this before opening the stream. An explicit header that
+      # cannot accept `text/event-stream` is a client that will not understand
+      # the response, so the host should answer 406 instead of streaming.
+      def self.accepts?(accept)
+        return true if accept.nil? || accept.strip.empty?
+
+        accept.split(',').any? do |range|
+          %w[text/event-stream text/* */*].include?(range.split(';').first.to_s.strip.downcase)
+        end
+      end
+
+      # The Content-Type for the response the host is about to open.
+      def content_type
+        @encoder.content_type
+      end
+
+      # The response headers a host should send before the first event.
+      def headers
+        self.class.headers(content_type)
+      end
+
+      # The response headers a host should send before the first event, for the
+      # default media type. A host that mounts a run from a controller merges
+      # these without having a stream yet.
+      def self.headers(content_type = CONTENT_TYPE)
+        {
+          'Content-Type' => "#{content_type}; charset=utf-8",
+          'Cache-Control' => CACHE_CONTROL,
+          'Connection' => CONNECTION,
+          'X-Accel-Buffering' => NO_TRANSFORM
+        }
+      end
+
+      # Encodes +event+ and writes it, then flushes so the client sees it now.
+      # Returns the encoded frame.
+      def write(event)
+        raise StreamError, 'The event stream is closed.' if @closed
+
+        frame = @encoder.encode(event)
+        @stream.write(frame)
+        flush
+        @events_written += 1
+        @bytes_written += frame.bytesize
+        frame
+      rescue IOError, SystemCallError => e
+        # The client is gone. The class name is enough for the host log; the
+        # message may embed a path or a socket address and is not echoed.
+        @logger&.warn("[ruby_llm_ag_ui] the event stream failed: #{e.class}")
+        @closed = true
+        raise StreamError, "The event stream failed: #{e.class}"
+      end
+
+      # Pushes buffered bytes to the client. A no-op for a stream that cannot
+      # flush, such as a Rack response body.
+      def flush
+        @stream.flush if @stream.respond_to?(:flush)
+      end
+
+      # Marks the stream closed. The host calls this once the run is over; a
+      # later write raises.
+      def close
+        @closed = true
+      end
+
+      # Whether the stream can no longer take events.
+      def closed?
+        @closed
+      end
+    end
+  end
+end

data/lib/ruby_llm/ag_ui/store/base.rb

@@ -0,0 +1,123 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module AgUI
+    # Run identity and progress storage.
+    #
+    # The store is where run identity becomes durable: it reserves a run before
+    # any provider work happens, so a retried request attaches to the run
+    # already in progress instead of executing twice, and it keeps the
+    # interrupts that a later resume has to answer.
+    #
+    # Two stores ship with the gem: Store::MemoryStore, which is the default for
+    # a single process, and the optional Rails adapter's ActiveRecord store. A
+    # host may implement this interface instead; every method must be safe to
+    # call from concurrent threads.
+    #
+    # `scope` is the host's own tenant/actor boundary. The store treats it as an
+    # opaque key, so the host must build it from authenticated identity (for
+    # example `"#{tenant.id}:#{actor.id}"`). Two scopes never share state, which
+    # is what stops one tenant from attaching to another's run.
+    module Store
+      # The interface every store implements. Subclass it to get the documented
+      # contract in one place, or duck-type it.
+      #
+      # The private helpers below are not part of the interface: they are the
+      # parts of the contract two stores would otherwise implement twice (the
+      # scope key, the clock, the replay rule). A store that duck-types the
+      # public methods has to satisfy them itself.
+      class Base
+        # Raised when a host passes a scope that cannot identify a tenant.
+        SCOPE_MESSAGE = 'scope must be a non-empty String built from authenticated identity'
+
+        def initialize(clock: nil)
+          @clock = clock || RubyLLM::AgUI.config.clock
+        end
+
+        # Claims (scope, thread_id, run_id) for +digest+.
+        #
+        # Returns a Reservation whose record is the run's authoritative state.
+        # A second call with the same key and an identical digest returns an
+        # ATTACHED reservation; a different digest raises RunConflictError,
+        # because the same run id may not mean two different requests.
+        def reserve(scope:, thread_id:, run_id:, digest:)
+          raise NotImplementedError
+        end
+
+        # The record for an exact identity, or nil.
+        def find(scope:, thread_id:, run_id:)
+          raise NotImplementedError
+        end
+
+        # The most recent run on the thread. Used to answer "what is this thread
+        # doing right now" without executing anything.
+        def latest_run(scope:, thread_id:)
+          raise NotImplementedError
+        end
+
+        # The most recent run on the thread that ended interrupted, or nil.
+        #
+        # A resume entry may only answer an interrupt of this run; anything
+        # older is stale and must be refused.
+        def latest_interrupted(scope:, thread_id:)
+          raise NotImplementedError
+        end
+
+        # Persists an updated record and returns it.
+        def save(record)
+          raise NotImplementedError
+        end
+
+        # The interrupt with +id+ anywhere on the thread, or nil.
+        def find_interrupt(scope:, thread_id:, interrupt_id:)
+          raise NotImplementedError
+        end
+
+        # Atomically records +decision+ on the interrupt with +interrupt_id+ and
+        # returns a Redemption.
+        #
+        # The first answer wins. A repeated answer with the same decision
+        # returns the stored interrupt as a replay, so a replayed resume request
+        # cannot run the tool twice; Redemption#recorded? says which call won.
+        # An expired interrupt, an unknown id, or a second answer that disagrees
+        # with the stored decision raises ResumeRejectedError: nothing is
+        # recorded and nothing is executed.
+        def redeem_interrupt(scope:, thread_id:, interrupt_id:, decision:, now:)
+          raise NotImplementedError
+        end
+
+        private
+
+        # The host's scope must be a canonical string. An arbitrary object would
+        # make two requests from the same tenant hash differently, so a retry
+        # would not attach and the run would execute twice.
+        def scope_key(scope)
+          raise ConfigurationError, SCOPE_MESSAGE unless scope.is_a?(String) && !scope.empty?
+
+          scope
+        end
+
+        def now
+          @clock.call
+        end
+
+        def blank?(value)
+          value.nil? || (value.respond_to?(:empty?) && value.empty?)
+        end
+
+        # A repeat of the stored decision is a replay: the caller is told the
+        # answer already given and must not execute the tool again. A second
+        # answer that disagrees rescinds nothing and is refused.
+        def replay(interrupt, decision)
+          return Redemption.new(interrupt: interrupt, replayed: true) if interrupt.status == decision
+
+          raise ResumeRejectedError, "interrupt #{interrupt.id} was already #{interrupt.status}"
+        end
+
+        def conflict_message(record)
+          "run #{record.run_id} on thread #{record.thread_id} was already reserved with a different input"
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/ag_ui/store/memory_store.rb

@@ -0,0 +1,123 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module AgUI
+    module Store
+      # An in-process run store.
+      #
+      # It is correct for a single process and a reasonable default in
+      # development and tests. It is not shared between processes, so a host
+      # that runs more than one Puma worker, or retries on a different worker,
+      # must use a durable store (the optional Rails adapter's ActiveRecord
+      # store) instead. Nothing here pretends otherwise.
+      #
+      # Every public method holds the store's lock for its whole body, so the
+      # reservation and the interrupt redemption are atomic: the first caller to
+      # answer an interrupt is the only one that records a decision. The private
+      # helpers assume the lock is already held and never take it themselves.
+      class MemoryStore < Base
+        def initialize(clock: nil)
+          super
+          @records = {}
+          @threads = {}
+          @lock = Mutex.new
+        end
+
+        def reserve(scope:, thread_id:, run_id:, digest:)
+          key = key_for(scope, thread_id, run_id)
+          @lock.synchronize do
+            existing = @records[key]
+            next attach(existing, digest) if existing
+
+            record = RunRecord.new(
+              scope: scope, thread_id: thread_id, run_id: run_id, digest: digest, created_at: now
+            )
+            @records[key] = record
+            thread = thread_key(scope, thread_id)
+            @threads[thread] = (@threads[thread] || []) + [run_id]
+            Reservation.new(record: record, status: Reservation::CREATED)
+          end
+        end
+
+        def find(scope:, thread_id:, run_id:)
+          key = key_for(scope, thread_id, run_id)
+          @lock.synchronize { @records[key] }
+        end
+
+        def latest_run(scope:, thread_id:)
+          @lock.synchronize { records_for(scope, thread_id).last }
+        end
+
+        def latest_interrupted(scope:, thread_id:)
+          @lock.synchronize { records_for(scope, thread_id).reverse.find(&:interrupted?) }
+        end
+
+        def save(record)
+          key = key_for(record.scope, record.thread_id, record.run_id)
+          @lock.synchronize do
+            stored = @records[key]
+            raise InvalidRunIdentityError, "run #{record.run_id} was never reserved" if stored.nil?
+            raise RunConflictError, conflict_message(record) if stored.digest != record.digest
+
+            @records[key] = record
+          end
+          record
+        end
+
+        def find_interrupt(scope:, thread_id:, interrupt_id:)
+          raise InvalidInputError, 'interrupt_id is required.' if blank?(interrupt_id)
+
+          @lock.synchronize { record_for_interrupt(scope, thread_id, interrupt_id)&.interrupt(interrupt_id) }
+        end
+
+        def redeem_interrupt(scope:, thread_id:, interrupt_id:, decision:, now:)
+          unless Interrupt::DECISIONS.include?(decision)
+            raise ConfigurationError, "decision must be one of #{Interrupt::DECISIONS.join(', ')}"
+          end
+
+          @lock.synchronize { redeem(scope, thread_id, interrupt_id, decision, now) }
+        end
+
+        private
+
+        def redeem(scope, thread_id, interrupt_id, decision, now)
+          record = record_for_interrupt(scope, thread_id, interrupt_id)
+          raise ResumeRejectedError, "no interrupt #{interrupt_id}" if record.nil?
+
+          interrupt = record.interrupt(interrupt_id)
+          return replay(interrupt, decision) if interrupt.decided?
+          raise ResumeRejectedError, "interrupt #{interrupt_id} expired" if interrupt.expired?(now)
+
+          answered = interrupt.decide(decision, now: now)
+          @records[key_for(scope, thread_id, record.run_id)] = record.replace_interrupt(answered, now: now)
+          Redemption.new(interrupt: answered, replayed: false)
+        end
+
+        # A repeat of the stored decision is a replay: the caller is told the
+        # answer already given and must not execute the tool again.
+        def record_for_interrupt(scope, thread_id, interrupt_id)
+          records_for(scope, thread_id).reverse.find { |record| record.interrupt(interrupt_id) }
+        end
+
+        def attach(existing, digest)
+          raise RunConflictError, conflict_message(existing) if existing.digest != digest
+
+          Reservation.new(record: existing, status: Reservation::ATTACHED)
+        end
+
+        def records_for(scope, thread_id)
+          thread = @threads[thread_key(scope, thread_id)] || []
+          thread.filter_map { |run_id| @records[key_for(scope, thread_id, run_id)] }
+        end
+
+        def key_for(scope, thread_id, run_id)
+          [scope_key(scope), thread_id, run_id]
+        end
+
+        def thread_key(scope, thread_id)
+          [scope_key(scope), thread_id]
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/ag_ui/validation.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module AgUI
+    # Shared shape checks for untrusted JSON.
+    #
+    # Every method raises an InputError subclass whose message names the
+    # offending member but never quotes the value, so a host can return it to
+    # the client without echoing request bytes back.
+    #
+    # Included by RubyLLM::AgUI::Input::Parser and RubyLLM::AgUI::MessageParser.
+    module Validation
+      # The longest accepted id (threadId, runId, message id, interrupt id).
+      MAX_ID_LENGTH = 255
+
+      # The longest accepted free-text string.
+      MAX_CONTENT_LENGTH = 1_000_000
+
+      private
+
+      def identifier!(value, name, required: true)
+        return nil if value.nil? && !required
+
+        raise InvalidInputError, "#{name} is required." if value.nil?
+        raise InvalidInputError, "#{name} must be a string." unless value.is_a?(String)
+        raise InvalidInputError, "#{name} must not be empty." if value.empty?
+        raise InvalidInputError, "#{name} is too long." if value.length > MAX_ID_LENGTH
+
+        value
+      end
+
+      def string!(value, name, max: MAX_CONTENT_LENGTH)
+        raise InvalidInputError, "#{name} is required." if value.nil?
+        raise InvalidInputError, "#{name} must be a string." unless value.is_a?(String)
+        raise InvalidInputError, "#{name} is too long." if value.length > max
+
+        value
+      end
+
+      def optional_string(value, name)
+        return nil if value.nil?
+
+        string!(value, name)
+      end
+
+      def array!(value, name, required: false, max: nil)
+        return [] if value.nil? && !required
+
+        raise InvalidInputError, "#{name} is required." if value.nil?
+        raise InvalidInputError, "#{name} must be an array." unless value.is_a?(Array)
+        raise InputTooLargeError, "#{name} exceeds the #{max} entry limit." if max && value.size > max
+
+        value
+      end
+
+      def object!(value, name)
+        raise InvalidInputError, "#{name} must be an object." unless value.is_a?(Hash)
+
+        value
+      end
+    end
+  end
+end

data/lib/ruby_llm/ag_ui/version.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module AgUI
+    # The version of the ruby_llm-ag_ui gem, as a string.
+    VERSION = '0.1.0'
+  end
+end

data/lib/ruby_llm/ag_ui.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+require 'ruby_llm'
+require 'ag_ui_protocol'
+require 'zeitwerk'
+require 'json'
+require 'digest'
+require 'securerandom'
+require 'time'
+
+require_relative 'ag_ui/version'
+
+# RubyLLM::AgUI is an unofficial RubyLLM ecosystem extension that runs one
+# RubyLLM 2.x chat turn as one AG-UI 1.0 run. Require it with:
+#
+#   require 'ruby_llm/ag_ui'
+#
+# The core is pure Ruby. Requiring this file does not load Rails, open a
+# network connection, or read provider credentials. The Rails adapter lives
+# under RubyLLM::AgUI::Rails and loads only when Rails is present.
+module RubyLLM
+  # The AG-UI integration namespace. Inside this namespace the protocol SDK is
+  # referred to as ::AgUiProtocol and the Rails framework as ::Rails, because
+  # RubyLLM::AgUI::Rails would otherwise shadow them.
+  module AgUI
+    class << self
+      # Yields the memoized configuration.
+      #
+      #   RubyLLM::AgUI.configure do |config|
+      #     config.run_store = MyRunStore.new
+      #   end
+      def configure
+        yield config
+      end
+
+      # The memoized RubyLLM::AgUI::Configuration.
+      def config
+        @config ||= Configuration.new
+      end
+
+      # Replaces the configuration with defaults. Intended for tests.
+      def reset_configuration!
+        @config = Configuration.new
+      end
+    end
+  end
+end
+
+loader = Zeitwerk::Loader.new
+loader.tag = 'ruby_llm-ag_ui'
+# Zeitwerk inflections match a whole basename, so each entry names a file:
+# sse_writer.rb -> SSEWriter. Acronyms inside a longer name (ui, api) are not
+# inflected and would need their own entry.
+loader.inflector.inflect('sse_writer' => 'SSEWriter')
+loader.push_dir(File.expand_path('ag_ui', __dir__), namespace: RubyLLM::AgUI)
+# Errors are part of the boot contract: hosts rescue them without a lazy load.
+loader.ignore(File.expand_path('ag_ui/error.rb', __dir__))
+# The Rails adapter is not part of the pure-Ruby core. It is loaded by its
+# railtie when Rails is present, or explicitly by the host.
+loader.ignore(File.expand_path('ag_ui/rails', __dir__))
+loader.setup
+
+require_relative 'ag_ui/error'
+
+require_relative 'ag_ui/rails/railtie' if defined?(Rails::Railtie)