ag-ui 0.2.0 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: ba6811e09735cfc95262ba711a5e953a92d9af8d78ea8b13867f1d1c60fbeedc
-  data.tar.gz: 75e995ce5c17476566258e05e5b56de4481b3df8539554e91d687363e15ad456
+  metadata.gz: 6092a88471f821fdd20d43bd3dad28205f114bb5b2068faea39234b4f1bc87da
+  data.tar.gz: 93566155eed8ad3464386a74b8d5afb6db6ea124eccac4bf6382b264e02733ec
 SHA512:
-  metadata.gz: 7e1bfaaa5c833b60ab7ce9a99d069abd3603eaf8366abf355b73688b07c54fb1d76fd280fb46a00e3e0e01a5cbf9de8d7b4ce7e59b581679eff92f7f891842c2
-  data.tar.gz: 205c4a2f519bcbfcf1fcb975de4929da568633d1358a793c78b042424c867918cce430dd6c9289b965ead31bb1a218a4d706b14e7476093a2ac42d7a10a1e21b
+  metadata.gz: cf108e6f97a63801cf5b3a936241d187a3492bf8c92209ff0353d5e98ae68bb3752c05aedeffcb8f7a1deee0d91378128dc9608b3a4f997e8a98b9712412fc89
+  data.tar.gz: 541b3da611c70fdd1ec02b97dd9c8fad76d30f20e124d38800896d57b431a099c148d3635193ebef28aaae6d70b303989e41380abeb6c4e489820cd2fc1d6d1c

data/lib/ag_ui/event_bridge.rb

@@ -24,12 +24,19 @@ module AgUi
       tool_call_args: :translate_tool_call_args,
       tool_call_end: :translate_tool_call_end,
       tool_call_result: :translate_tool_call_result,
+      state_snapshot: :translate_state_snapshot,
+      state_delta: :translate_state_delta,
+      messages_snapshot: :translate_messages_snapshot,
       activity_snapshot: :translate_activity_snapshot,
       reasoning_start: :translate_reasoning_start,
       reasoning_message_start: :translate_reasoning_message_start,
       reasoning_message_content: :translate_reasoning_message_content,
       reasoning_message_end: :translate_reasoning_message_end,
       reasoning_end: :translate_reasoning_end,
+      step_started: :translate_step_started,
+      step_finished: :translate_step_finished,
+      custom: :translate_custom,
+      raw: :translate_raw,
     }.freeze
 
     def initialize(stream)
@@ -88,6 +95,40 @@ module AgUi
         )
       end
 
+      # Shared state (CoAgents). snapshot replaces the whole state; delta is a
+      # JSON Patch (RFC 6902) op array the client applies to its own store.
+      def translate_state_snapshot(data)
+        @stream.state_snapshot(snapshot: data[:snapshot])
+      end
+
+      def translate_state_delta(data)
+        @stream.state_delta(delta: data[:delta])
+      end
+
+      def translate_messages_snapshot(data)
+        @stream.messages_snapshot(messages: data[:messages])
+      end
+
+      # Structured progress markers (paired start/finish by step name).
+      def translate_step_started(data)
+        @stream.step_started(step_name: data[:step_name])
+      end
+
+      def translate_step_finished(data)
+        @stream.step_finished(step_name: data[:step_name])
+      end
+
+      # Escape hatches: CUSTOM carries app/agent-defined events (e.g. the
+      # "PredictState" convention for predictive state updates); RAW passes a
+      # framework event straight through.
+      def translate_custom(data)
+        @stream.custom(name: data[:name], value: data[:value])
+      end
+
+      def translate_raw(data)
+        @stream.raw(event: data[:event], source: data[:source])
+      end
+
       def translate_activity_snapshot(data)
         @stream.activity_snapshot(
           message_id: data[:message_id],
@@ -171,4 +212,36 @@ describe "AgUi::EventBridge" do
     result.should.equal?(bridge)
     read_frames.(stream).should == []
   end
+
+  it "translates shared-state events (snapshot + JSON-Patch delta)" do
+    stream = AgUi::Server::SSE::Stream.new(thread_id: "t1", run_id: "r1")
+    bridge = AgUi::EventBridge.new(stream)
+
+    bridge << { type: :state_snapshot, data: { snapshot: { "theme" => "dark" } } }
+    bridge << { type: :state_delta,
+                data: { delta: [{ "op" => "replace", "path" => "/theme", "value" => "light" }] } }
+    stream.finish
+
+    frames = read_frames.(stream)
+    frames.map { |f| f["type"] }.should == %w[STATE_SNAPSHOT STATE_DELTA]
+    frames[0]["snapshot"].should == { "theme" => "dark" }
+    frames[1]["delta"].first["path"].should == "/theme"
+  end
+
+  it "translates CUSTOM (e.g. the PredictState convention) and STEP markers" do
+    stream = AgUi::Server::SSE::Stream.new(thread_id: "t1", run_id: "r1")
+    bridge = AgUi::EventBridge.new(stream)
+
+    bridge << { type: :step_started, data: { step_name: "plan" } }
+    bridge << { type: :custom,
+                data: { name: "PredictState",
+                        value: [{ "state_key" => "document", "tool" => "write" }] } }
+    bridge << { type: :step_finished, data: { step_name: "plan" } }
+    stream.finish
+
+    frames = read_frames.(stream)
+    frames.map { |f| f["type"] }.should == %w[STEP_STARTED CUSTOM STEP_FINISHED]
+    frames[0]["stepName"].should == "plan"
+    frames[1]["name"].should == "PredictState"
+  end
 end

data/lib/ag_ui/middleware/state.rb

@@ -0,0 +1,341 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "securerandom"
+require "hana"
+require "ag_ui"
+
+module AgUi
+  module Middleware
+    # Shared state (CoAgents) — the Ruby side of AG-UI's STATE_SNAPSHOT /
+    # STATE_DELTA channel (doc 05). Bidirectional state sync between the
+    # frontend's `agent.state` and the running agent.
+    #
+    # Way in:
+    #   - seeds env[:state] from RunAgentInput.state (the state the frontend
+    #     last held, pushed via agent.setState) so downstream middleware,
+    #     tools, and the app can READ what the UI currently shows
+    #   - injects two tools the model calls to WRITE state:
+    #       AGUISendStateSnapshot({ snapshot })  — replace the whole state
+    #       AGUISendStateDelta({ delta })        — RFC 6902 JSON Patch ops
+    #
+    # Way out, per state-tool call the model made:
+    #   - snapshot: env[:state] = snapshot; emits STATE_SNAPSHOT
+    #   - delta:    patches env[:state] (Hana); emits STATE_DELTA
+    #   - both append a {"status":"ok"} :tool message and emit a
+    #     TOOL_CALL_RESULT — the same shape a server tool produces — then let
+    #     the run CONTINUE (unlike a browser/client tool) so the model can act
+    #     on the new state or confirm to the user. State tools are server-side.
+    #
+    # The TOOL_CALL_START/ARGS/END chrome is left to ToolRouter (it advertises
+    # nothing about these being "client" vs "server" — it just streams the
+    # call), exactly as A2ui leaves its render_a2ui call to ToolRouter. Compose
+    # State OUTSIDE ToolRouter:
+    #
+    #   use Loop::ToolResult
+    #   use State, state: input.state
+    #   use ToolRouter, tools: input.tools
+    #
+    # The frontend applies the same snapshot/patch to its own store; env[:state]
+    # is kept coherent so the agent's own later reads (and further deltas) see
+    # the current value.
+    class State
+      SNAPSHOT_TOOL = "AGUISendStateSnapshot"
+      DELTA_TOOL    = "AGUISendStateDelta"
+      TOOL_NAMES    = [SNAPSHOT_TOOL, DELTA_TOOL].freeze
+
+      SNAPSHOT_DEFINITION = {
+        "name" => SNAPSHOT_TOOL,
+        "description" =>
+          "Replace the shared application state with a new snapshot; the " \
+          "frontend re-renders from it. Send the COMPLETE next state object, " \
+          "not a diff. Prefer AGUISendStateDelta for small targeted changes.",
+        "parameters" => {
+          "type" => "object",
+          "properties" => {
+            "snapshot" => {
+              "type" => "object",
+              "description" => "The complete new application state.",
+            },
+          },
+          "required" => %w[snapshot],
+        },
+      }.freeze
+
+      DELTA_DEFINITION = {
+        "name" => DELTA_TOOL,
+        "description" =>
+          "Apply a JSON Patch (RFC 6902) to the shared application state — an " \
+          "array of {op, path, value} operations. Use for small, targeted " \
+          "changes. Paths are JSON Pointers, e.g. " \
+          "\"/documentEditor/activeTabId\". ops: add, replace, remove, move, " \
+          "copy, test.",
+        "parameters" => {
+          "type" => "object",
+          "properties" => {
+            "delta" => {
+              "type" => "array",
+              "description" =>
+                "JSON Patch operations, e.g. " \
+                "[{\"op\":\"replace\",\"path\":\"/theme\",\"value\":\"dark\"}].",
+              "items" => { "type" => "object" },
+            },
+          },
+          "required" => %w[delta],
+        },
+      }.freeze
+
+      def initialize(app, state: nil)
+        @app = app
+        @initial_state = normalize(state)
+      end
+
+      def call(env)
+        seed_state(env)
+        advertise(env)
+
+        before = env[:messages].length
+        @app.call(env)
+
+        # Only this iteration's assistant message — seeded history can carry
+        # old state-tool turns that must not re-emit.
+        appended = env[:messages][before..] || []
+        assistant = appended.reverse.find { |m| m.respond_to?(:tool_call?) && m.tool_call? }
+        if assistant
+          apply_state_calls(env, assistant)
+        end
+
+        env
+      end
+
+      private
+
+        def apply_state_calls(env, assistant)
+          state_calls = assistant.tool_calls.select { |tc| TOOL_NAMES.include?(tc.name) }
+          unless state_calls.empty?
+            state_calls.each { |tool_call| handle(env, tool_call) }
+
+            # State tools don't end the run — they're server-side, not browser
+            # tools — so let Loop::ToolResult continue (last message is now
+            # :tool). Only override the exit for a PURE state turn; a real
+            # client tool mixed in still ends the run so the browser runs it.
+            if assistant.tool_calls.all? { |tc| TOOL_NAMES.include?(tc.name) }
+              env[:should_exit] = false
+            end
+          end
+        end
+
+        def seed_state(env)
+          unless env.key?(:state)
+            env[:state] = @initial_state
+          end
+        end
+
+        # Idempotent — the turn loop re-enters every iteration.
+        def advertise(env)
+          env[:tools] ||= []
+          names = env[:tools].map { |t| t.is_a?(Hash) ? t["name"] : nil }
+          [SNAPSHOT_DEFINITION, DELTA_DEFINITION].each do |definition|
+            unless names.include?(definition["name"])
+              env[:tools] << definition
+            end
+          end
+        end
+
+        def handle(env, tool_call)
+          case tool_call.name
+          when SNAPSHOT_TOOL then apply_snapshot(env, tool_call)
+          when DELTA_TOOL    then apply_delta(env, tool_call)
+          end
+        end
+
+        def apply_snapshot(env, tool_call)
+          snapshot = tool_call.arguments["snapshot"]
+          if snapshot.is_a?(Hash)
+            env[:state] = snapshot
+            env[:events] << { type: :state_snapshot, data: { snapshot: snapshot } }
+            ack(env, tool_call)
+          else
+            ack(env, tool_call, error: "snapshot must be an object")
+          end
+        end
+
+        def apply_delta(env, tool_call)
+          delta = tool_call.arguments["delta"]
+          if delta.is_a?(Array)
+            patch_state(env, tool_call, delta)
+          else
+            ack(env, tool_call, error: "delta must be an array of JSON Patch operations")
+          end
+        end
+
+        # Keep our copy coherent for later reads/deltas; the frontend applies
+        # the same patch to its own store. A bad patch is reported back to the
+        # model as a tool error rather than raising the run.
+        def patch_state(env, tool_call, delta)
+          patched = Hana::Patch.new(delta).apply(deep_dup(env[:state] || {}))
+        rescue StandardError => e
+          ack(env, tool_call, error: "invalid JSON Patch: #{e.message}")
+        else
+          env[:state] = patched
+          env[:events] << { type: :state_delta, data: { delta: delta } }
+          ack(env, tool_call)
+        end
+
+        # Append the tool result (so the assistant tool-call has its matching
+        # :tool message and Loop::ToolResult continues) and emit TOOL_CALL_RESULT
+        # on the wire — the exact shape server tools produce.
+        def ack(env, tool_call, error: nil)
+          content = error ? { "status" => "error", "error" => error } : { "status" => "ok" }
+          json = JSON.generate(content)
+          env[:messages].tool(json, tool_call_id: tool_call.id)
+          env[:events] << {
+            type: :tool_call_result,
+            data: {
+              message_id: SecureRandom.uuid,
+              tool_call_id: tool_call.id,
+              content: json,
+            },
+          }
+        end
+
+        def normalize(state)
+          if state.nil?
+            nil
+          elsif state.respond_to?(:to_h)
+            state.to_h
+          else
+            state
+          end
+        end
+
+        def deep_dup(obj)
+          case obj
+          when Hash  then obj.each_with_object({}) { |(k, v), h| h[k] = deep_dup(v) }
+          when Array then obj.map { |v| deep_dup(v) }
+          else obj
+          end
+        end
+    end
+  end
+end
+
+__END__
+
+describe "AgUi::Middleware::State" do
+  snapshot_call = ->(args) do
+    Brute::Message.new(
+      role: :assistant, content: nil,
+      tool_calls: [{ id: "tc1", name: "AGUISendStateSnapshot", arguments: args }],
+    )
+  end
+
+  delta_call = ->(args) do
+    Brute::Message.new(
+      role: :assistant, content: nil,
+      tool_calls: [{ id: "tc1", name: "AGUISendStateDelta", arguments: args }],
+    )
+  end
+
+  it "advertises both state tools idempotently and seeds env[:state] from input" do
+    seen = nil
+    mw = AgUi::Middleware::State.new(->(env) { seen = env }, state: { "theme" => "light" })
+    env = { messages: Brute.log, events: [], tools: [{ "name" => "navigate" }] }
+    mw.call(env)
+    mw.call(env) # second loop iteration — no dupes
+
+    seen[:tools].map { |t| t["name"] }.should ==
+      %w[navigate AGUISendStateSnapshot AGUISendStateDelta]
+    seen[:state].should == { "theme" => "light" }
+  end
+
+  it "seeds from a Definition-like state via to_h" do
+    definition = Object.new
+    def definition.to_h = { "count" => 1 }
+    seen = nil
+    AgUi::Middleware::State.new(->(env) { seen = env }, state: definition)
+      .call({ messages: Brute.log, events: [], tools: [] })
+    seen[:state].should == { "count" => 1 }
+  end
+
+  it "AGUISendStateSnapshot: sets state, emits STATE_SNAPSHOT + result, continues the run" do
+    terminal = ->(env) { env[:messages] << snapshot_call.("snapshot" => { "count" => 3 }) }
+    env = { messages: Brute.log, events: [], tools: [], should_exit: true }
+    AgUi::Middleware::State.new(terminal).call(env)
+
+    env[:state].should == { "count" => 3 }
+    env[:events].map { |e| e[:type] }.should == %i[state_snapshot tool_call_result]
+    env[:events][0][:data][:snapshot].should == { "count" => 3 }
+    env[:messages].last.role.should == :tool
+    env[:messages].last.tool_call_id.should == "tc1"
+    env[:messages].last.content.should == "{\"status\":\"ok\"}"
+    env[:should_exit].should == false
+  end
+
+  it "AGUISendStateDelta: patches state (Hana), emits STATE_DELTA with the raw ops" do
+    delta = [{ "op" => "replace", "path" => "/theme", "value" => "dark" }]
+    terminal = ->(env) { env[:messages] << delta_call.("delta" => delta) }
+    env = { messages: Brute.log, events: [], tools: [], state: { "theme" => "light" }, should_exit: true }
+    AgUi::Middleware::State.new(terminal).call(env)
+
+    env[:state].should == { "theme" => "dark" }
+    env[:events].map { |e| e[:type] }.should == %i[state_delta tool_call_result]
+    env[:events][0][:data][:delta].should == delta
+    env[:should_exit].should == false
+  end
+
+  it "adds a nested key via delta against seeded state" do
+    delta = [{ "op" => "add", "path" => "/documentEditor", "value" => { "activeTabId" => "doc-2" } }]
+    terminal = ->(env) { env[:messages] << delta_call.("delta" => delta) }
+    env = { messages: Brute.log, events: [], tools: [], should_exit: true }
+    AgUi::Middleware::State.new(terminal, state: { "documentEditor" => { "activeTabId" => "doc-1" } })
+      .call(env)
+
+    env[:state].should == { "documentEditor" => { "activeTabId" => "doc-2" } }
+    env[:events][0][:type].should == :state_delta
+  end
+
+  it "reports a bad patch back to the model as a tool error, no STATE_DELTA emitted" do
+    delta = [{ "op" => "replace", "path" => "/missing/deep", "value" => 1 }]
+    terminal = ->(env) { env[:messages] << delta_call.("delta" => delta) }
+    env = { messages: Brute.log, events: [], tools: [], state: {}, should_exit: true }
+    AgUi::Middleware::State.new(terminal).call(env)
+
+    env[:events].map { |e| e[:type] }.should == %i[tool_call_result]
+    JSON.parse(env[:events][0][:data][:content])["status"].should == "error"
+    env[:messages].last.role.should == :tool
+  end
+
+  it "rejects a non-array delta and a non-object snapshot as tool errors" do
+    [delta_call.("delta" => "nope"), snapshot_call.("snapshot" => "nope")].each do |msg|
+      terminal = ->(env) { env[:messages] << msg }
+      env = { messages: Brute.log, events: [], tools: [] }
+      AgUi::Middleware::State.new(terminal).call(env)
+      JSON.parse(env[:events].last[:data][:content])["status"].should == "error"
+    end
+  end
+
+  it "leaves should_exit set when a real client tool is mixed into the turn" do
+    terminal = ->(env) do
+      env[:messages] << Brute::Message.new(
+        role: :assistant, content: nil,
+        tool_calls: [
+          { id: "tc1", name: "AGUISendStateSnapshot", arguments: { "snapshot" => { "a" => 1 } } },
+          { id: "tc2", name: "navigate", arguments: { "path" => "/x" } },
+        ],
+      )
+    end
+    env = { messages: Brute.log, events: [], tools: [], should_exit: true }
+    AgUi::Middleware::State.new(terminal).call(env)
+
+    env[:events].any? { |e| e[:type] == :state_snapshot }.should == true
+    env[:should_exit].should == true # browser still needs to run navigate
+  end
+
+  it "does nothing on a plain text turn" do
+    env = { messages: Brute.log, events: [], tools: [] }
+    AgUi::Middleware::State.new(->(e) { e[:messages].assistant("hi") }).call(env)
+    env[:events].should == []
+    env.key?(:should_exit).should == false
+  end
+end

data/lib/ag_ui/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module AgUi
-  VERSION = "0.2.0"
+  VERSION = "0.3.0"
 end

data/lib/ag_ui.rb

@@ -25,6 +25,7 @@ require "ag_ui/event_bridge"
 require "ag_ui/middleware/system_prompt"
 require "ag_ui/middleware/forwarded_props"
 require "ag_ui/middleware/tool_router"
+require "ag_ui/middleware/state"
 require "ag_ui/a2ui/catalog"
 require "ag_ui/a2ui/validate"
 require "ag_ui/a2ui/recovery"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: ag-ui
 version: !ruby/object:Gem::Version
-  version: 0.2.0
+  version: 0.3.0
 platform: ruby
 authors:
 - AgUi Contributors
@@ -65,6 +65,20 @@ dependencies:
     - - "~>"
       - !ruby/object:Gem::Version
         version: '2.5'
+- !ruby/object:Gem::Dependency
+  name: hana
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.3'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.3'
 - !ruby/object:Gem::Dependency
   name: brute
   requirement: !ruby/object:Gem::Requirement
@@ -200,6 +214,7 @@ files:
 - lib/ag_ui/messages.rb
 - lib/ag_ui/middleware/a2ui.rb
 - lib/ag_ui/middleware/forwarded_props.rb
+- lib/ag_ui/middleware/state.rb
 - lib/ag_ui/middleware/system_prompt.rb
 - lib/ag_ui/middleware/tool_router.rb
 - lib/ag_ui/protocol/json_schema.rb