ruby_llm-ag_ui 0.1.0

data/lib/ruby_llm/ag_ui/protocol.rb

@@ -0,0 +1,130 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module AgUI
+    # The few AG-UI 1.0 wire shapes the community Ruby SDK does not yet model.
+    #
+    # The protocol SDK owns event models and wire encoding, and this gem reuses
+    # them for everything else. Two 1.0 shapes are missing from
+    # `ag-ui-protocol` 0.2.0:
+    #
+    # 1. `RunFinishedCancelledOutcome` — 0.2.0 has success and interrupt
+    #    outcomes only. A cancellation reported as success is exactly the bug
+    #    this gem exists to prevent, so the gem emits a schema-valid cancelled
+    #    outcome and offers the model upstream.
+    # 2. `RUN_STARTED.protocolVersion` — 1.0 lets a producer declare the
+    #    version it speaks so a consumer can see a downgrade. 0.2.0's
+    #    RunStartedEvent has no such field.
+    #
+    # Both classes subclass the SDK's own models, so the SDK encoder serializes
+    # them unchanged. They define no top-level `AgUi` constant.
+    module Protocol
+      # The AG-UI protocol version this producer speaks.
+      PROTOCOL_VERSION = '1.0'
+
+      # The pinned upstream commit of docs/spec/1.0/schema.json.
+      SCHEMA_COMMIT = '89cedd3fd549feeee41e50e1d7c6d373e549241a'
+
+      # RunFinishedOutcome with the 1.0 "cancelled" discriminator.
+      #
+      # A cancelled run was stopped before it completed and did not fail. It
+      # carries no interrupts, so the next run on the thread is an ordinary new
+      # run rather than a resume.
+      class CancelledOutcome < ::AgUiProtocol::Core::Types::Model
+        # The discriminator value the 1.0 schema requires.
+        TYPE = 'cancelled'
+
+        attr_reader :type
+
+        def initialize(type: nil)
+          raise ArgumentError, "CancelledOutcome.type must be #{TYPE.inspect}" if !type.nil? && type != TYPE
+
+          super()
+          @type = TYPE
+        end
+
+        def to_h
+          { type: @type }
+        end
+      end
+
+      # RunStartedEvent that declares the producer's protocol version.
+      class RunStartedEvent < ::AgUiProtocol::Core::Events::RunStartedEvent
+        attr_reader :protocol_version
+
+        def initialize(protocol_version: PROTOCOL_VERSION, **rest)
+          @protocol_version = protocol_version
+          super(**rest)
+        end
+
+        def to_h
+          super.merge(protocol_version: @protocol_version)
+        end
+      end
+
+      # RunFinishedEvent that accepts the cancelled outcome.
+      #
+      # The SDK's own RunFinishedEvent validates its +outcome+ against the two
+      # outcomes 0.2.0 models, so it cannot carry a cancelled one. This class
+      # has the same shape and the same RUN_FINISHED type; it accepts any
+      # outcome object that serializes to a schema-valid RunFinishedOutcome,
+      # including the SDK's success and interrupt outcomes.
+      class RunFinishedEvent < ::AgUiProtocol::Core::Events::BaseEvent
+        attr_reader :thread_id, :run_id, :result, :outcome, :usage
+
+        # Mirrors the SDK's own RunFinishedEvent signature.
+        def initialize(thread_id:, run_id:, result: nil, outcome: nil, usage: nil, timestamp: nil, raw_event: nil) # rubocop:disable Metrics/ParameterLists
+          if !result.nil? && !outcome.nil?
+            raise ArgumentError, 'result and outcome are mutually exclusive; provide one or the other, not both'
+          end
+
+          super(type: ::AgUiProtocol::Core::Events::EventType::RUN_FINISHED, timestamp: timestamp, raw_event: raw_event)
+          @thread_id = thread_id
+          @run_id = run_id
+          @result = result
+          @outcome = outcome
+          @usage = usage
+        end
+
+        def to_h
+          super.merge(thread_id: @thread_id, run_id: @run_id, result: @result, outcome: @outcome, usage: @usage)
+        end
+      end
+
+      module_function
+
+      # The RUN_STARTED this producer emits: the SDK's shape plus the protocol
+      # version only 1.0 carries.
+      def run_started(thread_id:, run_id:, parent_run_id: nil, input: nil)
+        RunStartedEvent.new(thread_id: thread_id, run_id: run_id, parent_run_id: parent_run_id, input: input)
+      end
+
+      # A RUN_FINISHED carrying +outcome+.
+      def run_finished(thread_id:, run_id:, outcome:)
+        RunFinishedEvent.new(thread_id: thread_id, run_id: run_id, outcome: outcome)
+      end
+
+      # The outcome of a run that ended without a human in the loop.
+      def success_outcome
+        ::AgUiProtocol::Core::Events::RunFinishedSuccessOutcome.new
+      end
+
+      # The 1.0 cancelled outcome 0.2.0 cannot express.
+      def cancelled_outcome
+        CancelledOutcome.new
+      end
+
+      # The outcome of a run that paused, carrying the interrupts that still
+      # wait for an answer. The schema requires at least one, so a run whose
+      # interrupts were all answered is not closed with this outcome.
+      def interrupt_outcome(interrupts)
+        ::AgUiProtocol::Core::Events::RunFinishedInterruptOutcome.new(interrupts: interrupts.map(&:to_protocol))
+      end
+
+      # The complete state a reconnect serves instead of replaying deltas.
+      def state_snapshot(state)
+        ::AgUiProtocol::Core::Events::StateSnapshotEvent.new(snapshot: state)
+      end
+    end
+  end
+end

data/lib/ruby_llm/ag_ui/rails/active_record_store.rb

@@ -0,0 +1,119 @@
+# frozen_string_literal: true
+
+require_relative 'record'
+
+module RubyLLM
+  module AgUI
+    module Rails
+      # The durable run store: run identity in a table instead of a process.
+      #
+      # The in-process store is correct for one process only. A host that runs
+      # more than one Puma worker, or that restarts, needs the reservation, the
+      # interrupts and the terminal status to outlive the worker that produced
+      # them, which is what this store is for.
+      #
+      # Ordering is by row id, so "the latest run on the thread" means the run
+      # reserved last, not the run whose clock reading is largest. Two requests
+      # that arrive in the same millisecond still have an order.
+      #
+      #     RubyLLM::AgUI.configure { |config| config.run_store = ActiveRecordStore.new }
+      class ActiveRecordStore < Store::Base
+        # Claims (scope, thread_id, run_id) for +digest+.
+        #
+        # The unique index on the identity is the real guard: two workers that
+        # insert at the same time cannot both win, and the loser attaches to the
+        # row the winner wrote instead of executing the turn a second time.
+        def reserve(scope:, thread_id:, run_id:, digest:)
+          scope_key(scope)
+          record = RunRecord.new(
+            scope: scope, thread_id: thread_id, run_id: run_id, digest: digest, created_at: now
+          )
+          Record.from_run_record(record).save!
+          Reservation.new(record: record, status: Reservation::CREATED)
+        rescue ::ActiveRecord::RecordNotUnique
+          attach(row_for(scope, thread_id, run_id), digest)
+        end
+
+        def find(scope:, thread_id:, run_id:)
+          row_for(scope, thread_id, run_id)&.as_run_record
+        end
+
+        def latest_run(scope:, thread_id:)
+          thread_rows(scope, thread_id).last&.as_run_record
+        end
+
+        def latest_interrupted(scope:, thread_id:)
+          interrupted = RunRecord::INTERRUPTED.to_s
+          thread_rows(scope, thread_id).reverse.find { |row| row.status == interrupted }&.as_run_record
+        end
+
+        def save(record)
+          scope_key(record.scope)
+          row = row_for(record.scope, record.thread_id, record.run_id)
+          raise InvalidRunIdentityError, "run #{record.run_id} was never reserved" if row.nil?
+          raise RunConflictError, conflict_message(record) if row.digest != record.digest
+
+          row.write_run_record(record)
+        end
+
+        def find_interrupt(scope:, thread_id:, interrupt_id:)
+          raise InvalidInputError, 'interrupt_id is required.' if blank?(interrupt_id)
+
+          row_for_interrupt(scope, thread_id, interrupt_id)&.as_run_record&.interrupt(interrupt_id)
+        end
+
+        # Records +decision+ on the interrupt under a row lock, so two resumes
+        # that arrive together cannot both record a decision and run the tool.
+        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
+
+          row = row_for_interrupt(scope, thread_id, interrupt_id)
+          raise ResumeRejectedError, "no interrupt #{interrupt_id}" if row.nil?
+
+          row.with_lock { redeem(row.reload.as_run_record, interrupt_id, decision, now) }
+        end
+
+        private
+
+        # A retry of the same identity attaches; a different input under the
+        # same key is a conflict, because one run id may not mean two requests.
+        def attach(row, digest)
+          raise RunConflictError, conflict_message(row.as_run_record) if row.digest != digest
+
+          Reservation.new(record: row.reload.as_run_record, status: Reservation::ATTACHED)
+        end
+
+        def redeem(record, interrupt_id, decision, now)
+          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)
+          row = row_for(record.scope, record.thread_id, record.run_id)
+          row.write_run_record(record.replace_interrupt(answered, now: now))
+          Redemption.new(interrupt: answered, replayed: false)
+        end
+
+        def row_for(scope, thread_id, run_id)
+          return nil if blank?(scope) || blank?(thread_id) || blank?(run_id)
+
+          Record.find_by(scope: scope, thread_id: thread_id, run_id: run_id)
+        end
+
+        def thread_rows(scope, thread_id)
+          Record.where(scope: scope, thread_id: thread_id).order(:id).to_a
+        end
+
+        # An interrupt id is unique among a thread's runs, so the thread is
+        # searched newest first and the first row that knows the id wins. A
+        # thread holds few runs, and the alternative is a JSON query that would
+        # bind this store to one database.
+        def row_for_interrupt(scope, thread_id, interrupt_id)
+          thread_rows(scope, thread_id).reverse.find { |row| row.as_run_record.interrupt(interrupt_id) }
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/ag_ui/rails/railtie.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module AgUI
+    module Rails
+      # Loads the optional Rails adapter, and nothing else.
+      #
+      # The core is pure Ruby and `require 'ruby_llm/ag_ui'` loads no Rails. This
+      # file is required only when Rails is already present, and it loads the
+      # adapter's parts in dependency order, skipping the Active Record store
+      # when the host has no Active Record at all. Calling it twice is harmless.
+      def self.setup!
+        require_relative 'streaming'
+        return unless defined?(::ActiveRecord::Base)
+
+        require_relative 'active_record_store'
+      end
+
+      # The railtie that loads the adapter during boot.
+      class Railtie < ::Rails::Railtie
+        initializer 'ruby_llm.ag_ui' do
+          RubyLLM::AgUI::Rails.setup!
+        end
+      end
+    end
+  end
+end
+
+RubyLLM::AgUI::Rails.setup!

data/lib/ruby_llm/ag_ui/rails/record.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+require 'active_record'
+
+module RubyLLM
+  module AgUI
+    module Rails
+      # The durable form of one core run record.
+      #
+      # The core RunRecord stays the authority; this row is where it lives
+      # between requests, so a run reserved on one Puma worker is the same run on
+      # the next. Identity is duplicated out of the payload into indexed columns
+      # so a thread can be listed without parsing JSON, and the payload is
+      # written whole because a record is immutable: there is no partial update
+      # to make.
+      #
+      # The payload is JSON *text*, not a database JSON column. The record is
+      # never queried by content, and encoding it here keeps the stored shape
+      # exactly `RunRecord#to_storage` on every database and every Rails version,
+      # instead of inheriting whichever encoder the host's ActiveSupport picks.
+      #
+      # Create the table with `bin/rails generate ruby_llm:ag_ui:install`.
+      class Record < ::ActiveRecord::Base
+        self.table_name = 'ruby_llm_ag_ui_runs'
+
+        # The core record this row stores.
+        def as_run_record
+          ::RubyLLM::AgUI::RunRecord.from_storage(::JSON.parse(payload))
+        end
+
+        # Writes +record+ onto this row and returns the record.
+        #
+        # The identity columns are not writable: a row keeps the identity it was
+        # reserved with, so a record cannot be stored under another run's key.
+        def write_run_record(record)
+          update!(digest: record.digest, status: record.status.to_s, sequence: record.sequence,
+                  payload: ::JSON.generate(record.to_storage))
+          record
+        end
+
+        # An unsaved row carrying +record+.
+        def self.from_run_record(record)
+          new(
+            scope: record.scope, thread_id: record.thread_id, run_id: record.run_id, digest: record.digest,
+            status: record.status.to_s, sequence: record.sequence, payload: ::JSON.generate(record.to_storage)
+          )
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/ag_ui/rails/streaming.rb

@@ -0,0 +1,166 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module AgUI
+    module Rails
+      # The controller concern that serves one AG-UI run as one SSE response.
+      #
+      #   class AgentRunsController < ApplicationController
+      #     include RubyLLM::AgUI::Rails::Streaming
+      #
+      #     def create
+      #       stream_ag_ui_run do |input|
+      #         chat = ag_ui_chats.find_by!(ag_ui_thread_id: input.thread_id)
+      #         authorize! chat
+      #         chat
+      #       end
+      #     end
+      #   end
+      #
+      # What the host owns, and the concern never guesses:
+      #
+      # * authentication (a `before_action`, or the action itself);
+      # * `#ag_ui_scope`, the tenant/actor key the store is bounded by — the
+      #   concern refuses to run without one, so an unauthenticated request
+      #   cannot reserve a run;
+      # * the chat the block returns, built with only the tools and the model
+      #   that actor may use. The client's `messages`, `tools` and `state` are
+      #   data for that chat; they authorize nothing.
+      #
+      # Refusals are decided before the response is committed, so a client that
+      # cannot read an event stream, or that sent an input the gem will not
+      # accept, gets a status code instead of half a run.
+      module Streaming
+        extend ::ActiveSupport::Concern
+
+        # A client that cannot read an event stream, and tells us so.
+        NOT_ACCEPTABLE = 406
+        # `resolve_chat` found no chat for the thread this actor owns.
+        NOT_FOUND = 404
+        # The run cannot be served: a live run already has a stream, or the
+        # transport is gone.
+        CONFLICT = 409
+
+        included do
+          include ::ActionController::Live
+        end
+
+        # Streams one run and returns the run's record, or nil when the client
+        # disconnected. The block receives the parsed input and returns the chat
+        # for it.
+        def stream_ag_ui_run(&resolve_chat)
+          if resolve_chat.nil?
+            raise ConfigurationError, 'stream_ag_ui_run needs a block that returns the chat for the input'
+          end
+
+          return reject_ag_ui_request(NOT_ACCEPTABLE, accept_message) unless ag_ui_accept?
+
+          input = ag_ui_input
+          chat = resolve_chat.call(input)
+          run = ag_ui_run(chat, input)
+
+          ag_ui_stream(run)
+        rescue ::RubyLLM::AgUI::InputError => e
+          # Every refusal the gem itself decides carries the status it should be
+          # answered with: 400 for an input that cannot be accepted, 409 for an
+          # identity claimed with different input or a resume that answers
+          # nothing, 413 for a body past the configured bound.
+          reject_ag_ui_request(e.http_status, e.message)
+        rescue ::ActiveRecord::RecordNotFound
+          reject_ag_ui_request(NOT_FOUND, 'There is no such thread.')
+        rescue ::RubyLLM::AgUI::StreamError => e
+          reject_ag_ui_request(CONFLICT, e.message)
+        end
+
+        # The tenant/actor key every reservation is bounded by.
+        #
+        # It must come from the authenticated request, not from the input: a
+        # client-chosen scope would let one actor attach to another's run.
+        def ag_ui_scope
+          raise ConfigurationError,
+                "#{self.class} must define #ag_ui_scope, built from the authenticated tenant and actor"
+        end
+
+        # The state to publish for a reconnect, or nil. The host's own run state
+        # is not the gem's to guess.
+        def ag_ui_state
+          nil
+        end
+
+        private
+
+        # Whether this client can read an event stream. An AG-UI client asks for
+        # `text/event-stream`; a request that asks for something else is told so
+        # instead of being streamed at.
+        def ag_ui_accept?
+          SSEWriter.accepts?(request.headers['Accept'])
+        end
+
+        def accept_message
+          "This endpoint streams #{SSEWriter::CONTENT_TYPE}. Send `Accept: #{SSEWriter::CONTENT_TYPE}`."
+        end
+
+        def ag_ui_input
+          ::RubyLLM::AgUI::Input.parse(ag_ui_body)
+        end
+
+        # The request body, bounded before it is parsed: an agent input is small,
+        # and a host must not read an unbounded body from a request into memory.
+        # The errors name the status a client should be answered with.
+        def ag_ui_body
+          limit = RubyLLM::AgUI.config.max_input_bytes
+          body = request.body.read(limit + 1)
+          raise MalformedInputError, 'the request body is empty.' if body.nil? || body.strip.empty?
+          raise InputTooLargeError, "the request body exceeds #{limit} bytes." if body.bytesize > limit
+
+          JSON.parse(body)
+        rescue JSON::ParserError
+          raise MalformedInputError, 'the request body is not JSON.'
+        end
+
+        def ag_ui_run(chat, input)
+          ::RubyLLM::AgUI::Run.new(
+            chat: chat, input: input, scope: ag_ui_scope, state: ag_ui_state
+          )
+        end
+
+        # Opens the SSE response, runs, and closes the stream whatever happens.
+        #
+        # A client that disconnects mid-run is not a run outcome: the record the
+        # store already holds says how the run ended, and a truncated stream says
+        # only that this client stopped reading.
+        def ag_ui_stream(run)
+          response.headers.merge!(SSEWriter.headers)
+          stream = response.stream
+          begin
+            run.stream(to: stream)
+          rescue ::ActionController::Live::ClientDisconnected, ::IOError
+            nil
+          rescue ::RubyLLM::AgUI::StreamError
+            # Before the first write this is a refusal the client can still be
+            # told about, so it goes back up. After it, the stream is already
+            # broken and the record in the store is the run's outcome: a
+            # truncated stream says only that this client stopped reading.
+            raise unless response.committed?
+
+            nil
+          ensure
+            stream.close
+          end
+        end
+
+        # Renders a refusal as a plain response. Everything refused here is
+        # refused before the run writes anything, so the client gets a status
+        # code rather than a stream that stops mid-event.
+        def reject_ag_ui_request(status, message)
+          if response.committed?
+            raise StreamError, 'the response is already committed; a refusal must be decided before the run starts'
+          end
+
+          response.headers['Content-Type'] = 'text/plain; charset=utf-8'
+          render plain: message, status: status
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/ag_ui/reconnect.rb

@@ -0,0 +1,86 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module AgUI
+    # What a client that re-sends an input whose run already ended gets (FR-05).
+    #
+    # A reconnect is not a second execution. The run ended, so the client gets
+    # its outcome and the state snapshot the run recorded instead of the deltas
+    # it may already have — or may have missed. Replaying deltas would apply text
+    # twice and could run a tool twice, so they are never replayed.
+    #
+    # What cannot be served this way is a run that has not ended: attaching a
+    # second stream to a live run needs a relay the gem does not own, and a
+    # silent second stream would duplicate the work. Those are refused.
+    class Reconnect
+      # The failure reported for a run that failed before its sanitized error
+      # was recorded — a record from an older release, or one whose write was
+      # interrupted.
+      UNRECORDED_FAILURE = {
+        'message' => 'The run failed before it could finish.', 'code' => 'internal_error'
+      }.freeze
+
+      # Whether +record+ can be served to a client that reconnects.
+      def self.serves?(record)
+        record.terminal?
+      end
+
+      def initialize(record:, input:)
+        @record = record
+        @input = input
+      end
+
+      # The events to write, in order: the run opens, declares the state it owns,
+      # and reports the outcome already recorded.
+      def events
+        [open_event, *snapshot_events, terminal_event]
+      end
+
+      private
+
+      def open_event
+        Protocol.run_started(
+          thread_id: @record.thread_id, run_id: @record.run_id,
+          parent_run_id: @input.parent_run_id, input: @input.to_h
+        )
+      end
+
+      def snapshot_events
+        return [] if @record.snapshot.nil?
+
+        [Protocol.state_snapshot(@record.snapshot)]
+      end
+
+      def terminal_event
+        case @record.status
+        when RunRecord::SUCCEEDED then finished(Protocol.success_outcome)
+        when RunRecord::CANCELLED then finished(Protocol.cancelled_outcome)
+        when RunRecord::INTERRUPTED then interrupted_event
+        else failure_event
+        end
+      end
+
+      def finished(outcome)
+        Protocol.run_finished(thread_id: @record.thread_id, run_id: @record.run_id, outcome: outcome)
+      end
+
+      # An interrupted run whose interrupts have all been answered was continued
+      # by a resume run, which owns what happened next. The run therefore has no
+      # interrupt to report (the schema requires at least one), and saying so is
+      # more honest than reporting answered interrupts as if they still waited.
+      def interrupted_event
+        return finished(Protocol.interrupt_outcome(@record.pending_interrupts)) if @record.pending_interrupts.any?
+
+        ::AgUiProtocol::Core::Events::RunErrorEvent.new(
+          message: 'This run ended interrupted, and every interrupt has been answered.',
+          code: 'run_resumed'
+        )
+      end
+
+      def failure_event
+        failure = @record.failure || UNRECORDED_FAILURE
+        ::AgUiProtocol::Core::Events::RunErrorEvent.new(message: failure['message'], code: failure['code'])
+      end
+    end
+  end
+end

data/lib/ruby_llm/ag_ui/redaction.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module AgUI
+    # The policy that decides what a tool result or activity payload may put on
+    # the wire.
+    #
+    # A tool result is arbitrary host data on its way to a browser. It may hold
+    # an API key, a customer record or an internal path, and the gem has no way
+    # to know which part is a secret. So the default policy publishes the shape
+    # and withholds the content, and a host that has decided its payloads are
+    # safe says so explicitly:
+    #
+    #   RubyLLM::AgUI.configure do |config|
+    #     config.redact = RubyLLM::AgUI::Redaction.method(:none)   # publish raw
+    #     config.redact = ->(kind, payload) { MyPolicy.call(kind, payload) }
+    #   end
+    #
+    # A pattern-matching scrubber is deliberately not the default: it would
+    # create confidence in something that cannot be correct, because the gem
+    # cannot know what the host considers personal.
+    module Redaction
+      # What a redacted value becomes.
+      PLACEHOLDER = '[redacted]'
+
+      # The kinds a policy is asked about.
+      TOOL_RESULT = :tool_result
+
+      module_function
+
+      # The default policy: keep the shape, withhold the content.
+      #
+      # A Hash keeps its keys, an Array keeps its length, and a scalar becomes
+      # the placeholder, so a payload still validates against the schema it
+      # travels in while carrying none of its values.
+      def default(_kind, payload)
+        case payload
+        when nil, '' then payload
+        when Hash then payload.to_h { |key, _value| [key, PLACEHOLDER] }
+        when Array then payload.map { PLACEHOLDER }
+        else PLACEHOLDER
+        end
+      end
+
+      # The opt-out: publish the payload unchanged. Use it only when the host
+      # has established that the payload is safe to show.
+      def none(_kind, payload)
+        payload
+      end
+
+      # Calls +callable+ with the arguments its arity takes, so a host can write
+      # either `->(kind, payload) { ... }` or the terse `->(payload) { ... }`.
+      def apply(callable, kind, payload)
+        callable.arity == 1 ? callable.call(payload) : callable.call(kind, payload)
+      end
+    end
+  end
+end