ag-ui 0.1.0

data/lib/ag_ui/run_store.rb

@@ -0,0 +1,180 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "async"
+require "async/queue"
+require "ag_ui"
+
+module AgUi
+  # Per-thread run bookkeeping backing /connect (replay + live attach) and
+  # /stop (cancel the run's Async task) — the Ruby counterpart of the Node
+  # runtime's InMemoryAgentRunner thread store.
+  #
+  # The interface is duck-typed; hosts can supply anything implementing
+  # begin_run / record / finish_run / attach_task / stop /
+  # open_subscription (e.g. a redis-backed store) via
+  # `AgUi.agent(store: ...)`.
+  module RunStore
+    # Single-process store. Fiber/thread-safe: state mutations are locked,
+    # and live fanout uses Async::Queue so connect streams receive events
+    # as they are recorded.
+    class InMemory
+      FINISHED = :finished
+
+      def initialize
+        @threads = Hash.new do |hash, key|
+          hash[key] = { historic: [], live: nil, task: nil, subscribers: [] }
+        end
+        @mutex = Mutex.new
+      end
+
+      def begin_run(thread_id, run_id)
+        @mutex.synchronize do
+          @threads[thread_id][:live] = { run_id: run_id, events: [] }
+        end
+      end
+
+      def record(thread_id, payload)
+        subscribers = @mutex.synchronize do
+          thread = @threads[thread_id]
+          if thread[:live]
+            thread[:live][:events] << payload
+            thread[:subscribers].dup
+          else
+            []
+          end
+        end
+        subscribers.each { |queue| queue.enqueue(payload) }
+      end
+
+      def finish_run(thread_id)
+        subscribers = @mutex.synchronize do
+          thread = @threads[thread_id]
+          if thread[:live]
+            thread[:historic] << thread[:live]
+            thread[:live] = nil
+          end
+          thread[:task] = nil
+          drained = thread[:subscribers]
+          thread[:subscribers] = []
+          drained
+        end
+        subscribers.each { |queue| queue.enqueue(FINISHED) }
+      end
+
+      # Only meaningful while the run is live — a late attach (the builder
+      # calls on_task after Async returns, which for an already-completed
+      # run is after finish_run) must not resurrect a stopped handle.
+      def attach_task(thread_id, task)
+        @mutex.synchronize do
+          if @threads[thread_id][:live]
+            @threads[thread_id][:task] = task
+          end
+        end
+      end
+
+      # Cancel the thread's in-flight run. The stream still closes cleanly
+      # (StreamBuilder's ensure) — matching the Node runtime, which ends
+      # an aborted run with a plain RUN_FINISHED-then-close.
+      def stop(thread_id)
+        task = @mutex.synchronize { @threads[thread_id][:task] }
+        if task
+          task.stop
+          true
+        else
+          false
+        end
+      end
+
+      # Atomic snapshot + live subscription: the returned events are
+      # everything recorded so far (historic runs + the live run's events),
+      # and when a run is in flight, queue receives each subsequent event
+      # and finally FINISHED. Atomicity means no gap and no duplicates
+      # between snapshot and subscription.
+      def open_subscription(thread_id)
+        @mutex.synchronize do
+          thread = @threads[thread_id]
+          events = thread[:historic].flat_map { |run| run[:events] }
+          queue = nil
+          if thread[:live]
+            events += thread[:live][:events]
+            queue = Async::Queue.new
+            thread[:subscribers] << queue
+          end
+          { events: events, queue: queue }
+        end
+      end
+    end
+  end
+end
+
+__END__
+
+describe "AgUi::RunStore::InMemory" do
+  it "replays historic runs in order for a thread" do
+    store = AgUi::RunStore::InMemory.new
+    store.begin_run("t1", "r1")
+    store.record("t1", { "type" => "RUN_STARTED" })
+    store.record("t1", { "type" => "RUN_FINISHED" })
+    store.finish_run("t1")
+    store.begin_run("t1", "r2")
+    store.record("t1", { "type" => "RUN_STARTED" })
+    store.finish_run("t1")
+
+    subscription = store.open_subscription("t1")
+    subscription[:events].map { |e| e["type"] }.should == %w[RUN_STARTED RUN_FINISHED RUN_STARTED]
+    subscription[:queue].should.be.nil
+  end
+
+  it "returns an empty snapshot for unknown threads" do
+    subscription = AgUi::RunStore::InMemory.new.open_subscription("nope")
+    subscription[:events].should == []
+    subscription[:queue].should.be.nil
+  end
+
+  it "fans live events out to subscribers, with FINISHED at run end" do
+    store = AgUi::RunStore::InMemory.new
+    store.begin_run("t1", "r1")
+    store.record("t1", { "type" => "RUN_STARTED" })
+
+    received = []
+    Async do |task|
+      subscription = store.open_subscription("t1")
+      subscription[:events].length.should == 1
+      queue = subscription[:queue]
+
+      reader = task.async do
+        loop do
+          item = queue.dequeue
+          if item == AgUi::RunStore::InMemory::FINISHED
+            break
+          end
+          received << item
+        end
+      end
+
+      store.record("t1", { "type" => "TEXT_MESSAGE_START" })
+      store.finish_run("t1")
+      reader.wait
+    end
+
+    received.map { |e| e["type"] }.should == ["TEXT_MESSAGE_START"]
+  end
+
+  it "stops the attached task exactly while a run is live" do
+    store = AgUi::RunStore::InMemory.new
+    stopped = false
+    fake_task = Object.new
+    fake_task.define_singleton_method(:stop) { stopped = true }
+
+    store.stop("t1").should == false
+
+    store.begin_run("t1", "r1")
+    store.attach_task("t1", fake_task)
+    store.stop("t1").should == true
+    stopped.should == true
+
+    store.finish_run("t1")
+    store.stop("t1").should == false
+  end
+end

data/lib/ag_ui/server/info.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "ag_ui"
+
+module AgUi
+  class Server
+    # Builds the GET /info payload — the envelope the CopilotKit client
+    # fetches on mount to discover agents and feature-detect capabilities.
+    # Shape extracted verbatim from @copilotkit/runtime's
+    # get-runtime-info.mjs (docs/09-ground-truth-host-app.md §1).
+    #
+    # The client feature-detects A2UI from the TOP-LEVEL `a2uiEnabled`
+    # flag; when enabled, the optional `a2ui: {enabled: true}` detail
+    # object rides along.
+    module Info
+      # The runtime version the client was built against; advertise the
+      # same so version-gated client behaviour matches the Node sidecar.
+      VERSION_PARITY = "1.62.2"
+
+      class << self
+        def payload(agent_id: "default", description: nil, a2ui_enabled: false, overrides: {})
+          agent = { "name" => agent_id, "className" => "BuiltInAgent" }
+          unless description.nil?
+            agent["description"] = description
+          end
+
+          base = {
+            "version" => VERSION_PARITY,
+            "agents" => { agent_id => agent },
+            "audioFileTranscriptionEnabled" => false,
+            "mode" => "sse",
+            "threadEndpoints" => {
+              "list" => false,
+              "inspect" => false,
+              "mutations" => false,
+              "realtimeMetadata" => false,
+            },
+            "a2uiEnabled" => a2ui_enabled,
+            "openGenerativeUIEnabled" => false,
+            "telemetryDisabled" => true,
+          }
+
+          if a2ui_enabled
+            base["a2ui"] = { "enabled" => true }
+          end
+
+          base.merge(overrides)
+        end
+      end
+    end
+  end
+end
+
+__END__
+
+describe "AgUi::Server::Info" do
+  it "matches the Node runtime envelope for the default agent" do
+    payload = AgUi::Server::Info.payload
+
+    payload["version"].should == "1.62.2"
+    payload["agents"].should == { "default" => { "name" => "default", "className" => "BuiltInAgent" } }
+    payload["mode"].should == "sse"
+    payload["a2uiEnabled"].should == false
+    payload.key?("a2ui").should == false
+    payload["threadEndpoints"].should == {
+      "list" => false, "inspect" => false, "mutations" => false, "realtimeMetadata" => false,
+    }
+  end
+
+  it "advertises a2ui via the top-level flag plus detail object" do
+    payload = AgUi::Server::Info.payload(a2ui_enabled: true)
+    payload["a2uiEnabled"].should == true
+    payload["a2ui"].should == { "enabled" => true }
+  end
+
+  it "includes agent description when given and applies overrides" do
+    payload = AgUi::Server::Info.payload(
+      agent_id: "host-app", description: "Studio assistant",
+      overrides: { "telemetryDisabled" => false },
+    )
+    payload["agents"]["host-app"]["description"].should == "Studio assistant"
+    payload["telemetryDisabled"].should == false
+  end
+end

data/lib/ag_ui/server/middleware/sse_stream.rb

@@ -0,0 +1,160 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "ag_ui"
+require "async"
+
+module AgUi
+  class Server
+    module Middleware
+      # Sets up an SSE stream builder on `env["ag_ui.stream"]`, ported
+      # from a2a's Server::Middleware::SSEStream.
+      #
+      # The `open` block runs inside an Async fiber and the stream is
+      # automatically finished when the block exits (even on exception).
+      # If the handler never calls `open`, the builder is removed from env
+      # so the response layer doesn't mistake it for a real stream.
+      #
+      # Usage (inside the /run handler):
+      #
+      #   env["ag_ui.stream"].open(thread_id:, run_id:) do |s|
+      #     s.run_started
+      #     s.text_message_content(message_id: "m1", delta: "Hi")
+      #     s.run_finished
+      #   end
+      #
+      class SSEStream
+        def initialize(app)
+          @app = app
+        end
+
+        def call(env)
+          builder = StreamBuilder.new(env)
+          env["ag_ui.stream"] = builder
+
+          result = @app.call(env)
+
+          # If open was never called, clear the builder so the response
+          # layer doesn't mistake it for a real stream.
+          if env["ag_ui.stream"].equal?(builder)
+            env.delete("ag_ui.stream")
+          end
+
+          result
+        end
+      end
+
+      # Factory that creates the SSE stream and runs the caller's block
+      # inside Async with automatic finish on exit.
+      #
+      # Created by SSEStream middleware — not intended for direct use.
+      class StreamBuilder
+        def initialize(env)
+          @env = env
+        end
+
+        # Create and open the SSE stream for the current run.
+        #
+        # The block runs inside an Async fiber; the stream is finished
+        # when the block exits, even if an exception is raised. The run
+        # handler owns terminal-event semantics (RUN_FINISHED / RUN_ERROR)
+        # — this layer only guarantees the connection closes.
+        #
+        # on_event / on_finish / on_task are the run-store taps (set by
+        # the Server's recording wrapper): every payload, end-of-run, and
+        # the Async task handle for /stop cancellation.
+        def open(thread_id:, run_id:, validate: true,
+                 on_event: nil, on_finish: nil, on_task: nil, &block)
+          stream = SSE::Stream.new(
+            thread_id: thread_id,
+            run_id: run_id,
+            validate: validate,
+            on_event: on_event,
+          )
+
+          @env["ag_ui.stream"] = stream
+
+          task = Async do
+            block.call(stream)
+          ensure
+            stream.finish
+            on_finish&.call
+          end
+          on_task&.call(task)
+
+          nil
+        end
+      end
+    end
+  end
+end
+
+__END__
+
+describe "AgUi::Server::Middleware::SSEStream" do
+  it "sets a StreamBuilder on env and clears it when open is never called" do
+    seen = nil
+    mw = AgUi::Server::Middleware::SSEStream.new(->(env) { seen = env["ag_ui.stream"]; :ok })
+
+    env = {}
+    result = mw.call(env)
+
+    seen.should.be.kind_of(AgUi::Server::Middleware::StreamBuilder)
+    env.key?("ag_ui.stream").should == false
+    result.should == :ok
+  end
+
+  it "preserves the opened stream on env" do
+    mw = AgUi::Server::Middleware::SSEStream.new(->(env) do
+      env["ag_ui.stream"].open(thread_id: "t1", run_id: "r1") { |s| s.run_started }
+    end)
+
+    env = {}
+    mw.call(env)
+
+    env["ag_ui.stream"].should.be.kind_of(AgUi::Server::SSE::Stream)
+  end
+end
+
+describe "AgUi::Server::Middleware::StreamBuilder" do
+  it "passes thread_id and run_id through to the stream" do
+    env = {}
+    builder = AgUi::Server::Middleware::StreamBuilder.new(env)
+
+    builder.open(thread_id: "t1", run_id: "r1") do |s|
+      s.thread_id.should == "t1"
+      s.run_id.should == "r1"
+    end
+  end
+
+  it "auto-finishes the stream when the block completes" do
+    env = {}
+    builder = AgUi::Server::Middleware::StreamBuilder.new(env)
+
+    builder.open(thread_id: "t1", run_id: "r1") { |s| s.run_started }
+
+    stream = env["ag_ui.stream"]
+    stream.read.should.include?("RUN_STARTED")
+    stream.read.should.be.nil
+  end
+
+  it "auto-finishes even when the block raises" do
+    env = {}
+    builder = AgUi::Server::Middleware::StreamBuilder.new(env)
+
+    builder.open(thread_id: "t1", run_id: "r1") do |s|
+      s.run_started
+      raise "boom"
+    end
+
+    stream = env["ag_ui.stream"]
+    stream.read.should.include?("RUN_STARTED")
+    stream.read.should.be.nil
+  end
+
+  it "returns nil" do
+    env = {}
+    builder = AgUi::Server::Middleware::StreamBuilder.new(env)
+    builder.open(thread_id: "t1", run_id: "r1") { |_s| }.should.be.nil
+  end
+end

data/lib/ag_ui/server/sse/event_encoder.rb

@@ -0,0 +1,69 @@
+# frozen_string_literal: true
+
+require "json"
+
+module AgUi
+  class Server
+    module SSE
+      # AG-UI SSE framing, ported from the reference encoder
+      # (ag-ui sdks/python/ag_ui/encoder/encoder.py, cross-checked against
+      # @ag-ui/encoder dist): one event per frame, `data: <json>\n\n`, no
+      # `event:`/`id:` lines, no heartbeats. Keys are camelCase and nil
+      # fields are omitted at every depth (pydantic exclude_none).
+      class EventEncoder
+        CONTENT_TYPE = "text/event-stream"
+
+        def content_type = CONTENT_TYPE
+
+        # event: a Hash with camelCase string/symbol keys, e.g.
+        #   { type: "TEXT_MESSAGE_CONTENT", messageId: "m1", delta: "Hi" }
+        def encode(event)
+          "data: #{JSON.generate(strip_nils(event))}\n\n"
+        end
+
+        private
+
+        def strip_nils(value)
+          case value
+          when Hash
+            value.each_with_object({}) do |(k, v), out|
+              unless v.nil?
+                out[k] = strip_nils(v)
+              end
+            end
+          when Array
+            value.map { |v| strip_nils(v) }
+          else
+            value
+          end
+        end
+      end
+    end
+  end
+end
+
+__END__
+
+describe "ag_ui/server/sse/event_encoder" do
+  encoder = AgUi::Server::SSE::EventEncoder.new
+
+  it "frames an event as a single data line with a blank-line terminator" do
+    frame = encoder.encode({ type: "TEXT_MESSAGE_CONTENT", messageId: "m1", delta: "Hi" })
+    frame.should == "data: {\"type\":\"TEXT_MESSAGE_CONTENT\",\"messageId\":\"m1\",\"delta\":\"Hi\"}\n\n"
+  end
+
+  it "omits nil fields at every depth" do
+    frame = encoder.encode({ type: "RUN_FINISHED", threadId: "t", runId: "r",
+                             result: nil, nested: { keep: 1, drop: nil } })
+    frame.should == "data: {\"type\":\"RUN_FINISHED\",\"threadId\":\"t\",\"runId\":\"r\",\"nested\":{\"keep\":1}}\n\n"
+  end
+
+  it "preserves nils inside arrays and keeps empty strings/false" do
+    frame = encoder.encode({ type: "CUSTOM", name: "x", value: ["a", nil], flag: false, s: "" })
+    frame.should == "data: {\"type\":\"CUSTOM\",\"name\":\"x\",\"value\":[\"a\",null],\"flag\":false,\"s\":\"\"}\n\n"
+  end
+
+  it "advertises the SSE content type" do
+    encoder.content_type.should == "text/event-stream"
+  end
+end