ag-ui 0.2.0 → 1.0.0

data/lib/ag_ui/middleware/forwarded_props.rb

@@ -9,7 +9,7 @@ module AgUi
     # middleware and the terminal can honor client-driven knobs — most
     # importantly `toolChoice` ({type: "function", function: {name}}),
     # which the suggestions engine uses to force `copilotkitSuggest`.
-    class ForwardedProps
+    class ForwardedProps < Brute::Middleware::Base
       def initialize(app, props: nil)
         @app = app
         @props = props

data/lib/ag_ui/middleware/state.rb

@@ -0,0 +1,366 @@
+# 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 < Brute::Middleware::Base
+      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.emit(:state_snapshot, { 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.emit(:state_delta, { 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)
+          if error
+            content = { "status" => "error", "error" => error }
+          else
+            content = { "status" => "ok" }
+          end
+          json = JSON.generate(content)
+          env[:messages].tool(json, tool_call_id: tool_call.id)
+          env.emit(
+            :tool_call_result,
+            {
+              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
+
+  # Brute 6 replaced the env[:events] sink with the hooks registry, so a test
+  # that wants to see what a middleware emitted subscribes instead of reading
+  # an array off the env. This records the vocabulary State speaks, in
+  # emission order.
+  STATE_EVENTS = %i[state_snapshot state_delta tool_call_result].freeze
+
+  def recording(env, into)
+    hooks = Brute::Hooks::Registry.new
+    STATE_EVENTS.each do |event|
+      hooks.on(event) { |_env, data, _trace| into << { type: event, data: data } }
+    end
+    Brute::Hooks::Trace.new(env, hooks: hooks)
+  end
+  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, tools: [{ "name" => "navigate" }] }
+    mw.call(recording(env, []))
+    mw.call(recording(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(recording({ messages: Brute.log, 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, tools: [], should_exit: true }
+    events = []
+    AgUi::Middleware::State.new(terminal).call(recording(env, events))
+
+    env[:state].should == { "count" => 3 }
+    events.map { |e| e[:type] }.should == %i[state_snapshot tool_call_result]
+    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, tools: [], state: { "theme" => "light" }, should_exit: true }
+    events = []
+    AgUi::Middleware::State.new(terminal).call(recording(env, events))
+
+    env[:state].should == { "theme" => "dark" }
+    events.map { |e| e[:type] }.should == %i[state_delta tool_call_result]
+    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, tools: [], should_exit: true }
+    events = []
+    AgUi::Middleware::State.new(terminal, state: { "documentEditor" => { "activeTabId" => "doc-1" } })
+      .call(recording(env, events))
+
+    env[:state].should == { "documentEditor" => { "activeTabId" => "doc-2" } }
+    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, tools: [], state: {}, should_exit: true }
+    events = []
+    AgUi::Middleware::State.new(terminal).call(recording(env, events))
+
+    events.map { |e| e[:type] }.should == %i[tool_call_result]
+    JSON.parse(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, tools: [] }
+      events = []
+      AgUi::Middleware::State.new(terminal).call(recording(env, events))
+      JSON.parse(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, tools: [], should_exit: true }
+    events = []
+    AgUi::Middleware::State.new(terminal).call(recording(env, events))
+
+    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, tools: [] }
+    events = []
+    AgUi::Middleware::State.new(->(e) { e[:messages].assistant("hi") }).call(recording(env, events))
+    events.should == []
+    env.key?(:should_exit).should == false
+  end
+end

data/lib/ag_ui/middleware/system_prompt.rb

@@ -11,7 +11,7 @@ module AgUi
     #
     # Skips entirely when the history already carries a system message
     # (the client can send its own via system/developer roles).
-    class SystemPrompt
+    class SystemPrompt < Brute::Middleware::Base
       def initialize(app, prompt: nil, context: nil)
         @app = app
         @prompt = prompt

data/lib/ag_ui/middleware/tool_router.rb

@@ -26,7 +26,7 @@ module AgUi
     #     RUN_FINISHED and the browser executes it (multi-run model)
     #   - mixed turns: server tools still execute, but any client call
     #     ends the run
-    class ToolRouter
+    class ToolRouter < Brute::Middleware::Base
       def initialize(app, tools: nil, server_tools: nil)
         @app = app
         @client_tools = tools || []
@@ -72,9 +72,9 @@ module AgUi
 
         def server_definition(tool)
           {
-            "name" => tool[:name].to_s,
+            "name"        => tool[:name].to_s,
             "description" => tool[:description].to_s,
-            "parameters" => tool[:parameters] || { "type" => "object" },
+            "parameters"  => tool[:parameters] || { "type" => "object" },
           }
         end
 
@@ -82,7 +82,7 @@ module AgUi
           client_called = false
 
           tool_calls.each do |tool_call|
-            emit_call(env[:events], tool_call)
+            emit_call(env, tool_call)
 
             server = @server_tools[tool_call.name]
             if server
@@ -97,19 +97,16 @@ module AgUi
           end
         end
 
-        def emit_call(events, tool_call)
-          events << {
-            type: :tool_call_start,
-            data: { tool_call_id: tool_call.id, tool_call_name: tool_call.name },
-          }
-          events << {
-            type: :tool_call_args,
-            data: { tool_call_id: tool_call.id, delta: JSON.generate(tool_call.arguments) },
-          }
-          events << {
-            type: :tool_call_end,
-            data: { tool_call_id: tool_call.id },
-          }
+        def emit_call(env, tool_call)
+          env.emit(
+            :tool_call_start,
+            { tool_call_id: tool_call.id, tool_call_name: tool_call.name },
+          )
+          env.emit(
+            :tool_call_args,
+            { tool_call_id: tool_call.id, delta: JSON.generate(tool_call.arguments) },
+          )
+          env.emit(:tool_call_end, { tool_call_id: tool_call.id })
         end
 
         def execute_server_tool(env, tool_call, tool)
@@ -120,16 +117,20 @@ module AgUi
             result = { "error" => e.message }
           end
 
-          content = result.is_a?(String) ? result : JSON.generate(result)
+          if result.is_a?(String)
+            content = result
+          else
+            content = JSON.generate(result)
+          end
           env[:messages].tool(content, tool_call_id: tool_call.id)
-          env[:events] << {
-            type: :tool_call_result,
-            data: {
-              message_id: SecureRandom.uuid,
+          env.emit(
+            :tool_call_result,
+            {
+              message_id:   SecureRandom.uuid,
               tool_call_id: tool_call.id,
-              content: content,
+              content:      content,
             },
-          }
+          )
         end
     end
   end
@@ -138,6 +139,20 @@ end
 __END__
 
 describe "AgUi::Middleware::ToolRouter" do
+  # Brute 6 replaced the env[:events] sink with the hooks registry, so a test
+  # that wants to see what a middleware emitted subscribes instead of reading
+  # an array off the env. This records the AG-UI tool vocabulary in emission
+  # order, which keeps the assertions below the shape they always had.
+  TOOL_EVENTS = %i[tool_call_start tool_call_args tool_call_end tool_call_result].freeze
+
+  def recording(env, into)
+    hooks = Brute::Hooks::Registry.new
+    TOOL_EVENTS.each do |event|
+      hooks.on(event) { |_env, data, _trace| into << { type: event, data: data } }
+    end
+    Brute::Hooks::Trace.new(env, hooks: hooks)
+  end
+
   it "advertises client and server tools idempotently across iterations" do
     seen = nil
     server_tool = { name: "get_time", description: "Now", handler: -> (_args) { "12:00" } }
@@ -147,9 +162,9 @@ describe "AgUi::Middleware::ToolRouter" do
       server_tools: [server_tool],
     )
 
-    env = { messages: Brute.log, events: [] }
-    mw.call(env)
-    mw.call(env) # second loop iteration
+    env = { messages: Brute.log }
+    mw.call(recording(env, []))
+    mw.call(recording(env, [])) # second loop iteration
 
     seen.map { |t| t["name"] }.should == %w[navigate get_time]
     seen.last["parameters"].should == { "type" => "object" }
@@ -169,17 +184,18 @@ describe "AgUi::Middleware::ToolRouter" do
       )
     end
 
-    env = { messages: Brute.log, events: [] }
-    AgUi::Middleware::ToolRouter.new(terminal, server_tools: [server_tool]).call(env)
+    env = { messages: Brute.log }
+    events = []
+    AgUi::Middleware::ToolRouter.new(terminal, server_tools: [server_tool]).call(recording(env, events))
 
     env[:messages].last.role.should == :tool
     env[:messages].last.tool_call_id.should == "tc1"
     env[:messages].last.content.should == "{\"found\":42}"
 
-    env[:events].map { |e| e[:type] }.should == %i[
+    events.map { |e| e[:type] }.should == %i[
       tool_call_start tool_call_args tool_call_end tool_call_result
     ]
-    env[:events].last[:data][:content].should == "{\"found\":42}"
+    events.last[:data][:content].should == "{\"found\":42}"
     env.key?(:should_exit).should == false
   end
 
@@ -192,8 +208,8 @@ describe "AgUi::Middleware::ToolRouter" do
       )
     end
 
-    env = { messages: Brute.log, events: [] }
-    AgUi::Middleware::ToolRouter.new(terminal, server_tools: [server_tool]).call(env)
+    env = { messages: Brute.log }
+    AgUi::Middleware::ToolRouter.new(terminal, server_tools: [server_tool]).call(recording(env, []))
 
     env[:messages].last.content.should == "{\"error\":\"kaput\"}"
     env.key?(:should_exit).should == false
@@ -211,11 +227,12 @@ describe "AgUi::Middleware::ToolRouter" do
       )
     end
 
-    env = { messages: Brute.log, events: [] }
-    AgUi::Middleware::ToolRouter.new(terminal, server_tools: [server_tool]).call(env)
+    env = { messages: Brute.log }
+    events = []
+    AgUi::Middleware::ToolRouter.new(terminal, server_tools: [server_tool]).call(recording(env, events))
 
     env[:should_exit].should == true
-    env[:events].count { |e| e[:type] == :tool_call_result }.should == 1
+    events.count { |e| e[:type] == :tool_call_result }.should == 1
   end
 
   it "emits TOOL_CALL events and exits when the turn ends on client tool calls" do
@@ -229,23 +246,25 @@ describe "AgUi::Middleware::ToolRouter" do
       )
     end
 
-    env = { messages: Brute.log, events: [] }
-    AgUi::Middleware::ToolRouter.new(terminal).call(env)
+    env = { messages: Brute.log }
+    events = []
+    AgUi::Middleware::ToolRouter.new(terminal).call(recording(env, events))
 
-    env[:events].map { |e| e[:type] }.should == %i[
+    events.map { |e| e[:type] }.should == %i[
       tool_call_start tool_call_args tool_call_end
       tool_call_start tool_call_args tool_call_end
     ]
-    env[:events][0][:data].should == { tool_call_id: "tc1", tool_call_name: "navigate" }
-    env[:events][1][:data][:delta].should == "{\"path\":\"/data\"}"
+    events[0][:data].should == { tool_call_id: "tc1", tool_call_name: "navigate" }
+    events[1][:data][:delta].should == "{\"path\":\"/data\"}"
     env[:should_exit].should == true
   end
 
   it "does nothing on the way out for plain text turns" do
-    env = { messages: Brute.log, events: [] }
-    AgUi::Middleware::ToolRouter.new(->(e) { e[:messages].assistant("hi") }).call(env)
+    env = { messages: Brute.log }
+    events = []
+    AgUi::Middleware::ToolRouter.new(->(e) { e[:messages].assistant("hi") }).call(recording(env, events))
 
-    env[:events].should == []
+    events.should == []
     env.key?(:should_exit).should == false
   end
 end

data/lib/ag_ui/protocol/json_schema/definition.rb

@@ -33,12 +33,12 @@ module AgUi
             camel = snake[k] || k
 
             if props.include?(camel)
-              @data[camel] = if value.is_a?(Definition)
-                value.to_h
+              if value.is_a?(Definition)
+                @data[camel] = value.to_h
               elsif (ref_info = refs[camel])
-                wrap_ref(value, ref_info)
+                @data[camel] = wrap_ref(value, ref_info)
               else
-                value
+                @data[camel] = value
               end
             end
           end
@@ -91,7 +91,7 @@ module AgUi
             raise ValidationError.new(
               errors,
               definition_name: self.class.definition_name,
-              data: to_h,
+              data:            to_h,
             )
           end
         end

data/lib/ag_ui/protocol/json_schema.rb

@@ -129,14 +129,18 @@ module AgUi
             properties.each do |camel_key, prop_schema|
               schema = unwrap_optional(prop_schema)
 
-              kind, ref =
-                if (r = schema["$ref"])
-                  [:object, r]
-                elsif schema["type"] == "array" && (r = schema.dig("items", "$ref"))
-                  [:array, r]
-                elsif schema["type"] == "object" && (r = schema.dig("additionalProperties", "$ref"))
-                  [:map, r]
-                end
+              kind = nil
+              ref = nil
+              if (r = schema["$ref"])
+                kind = :object
+                ref = r
+              elsif schema["type"] == "array" && (r = schema.dig("items", "$ref"))
+                kind = :array
+                ref = r
+              elsif schema["type"] == "object" && (r = schema.dig("additionalProperties", "$ref"))
+                kind = :map
+                ref = r
+              end
 
               if ref
                 name = ref_name_for(ref)
@@ -168,13 +172,13 @@ module AgUi
           end
 
           def build_snake_to_camel(camel_keys)
-            map = {}
-            camel_keys.each do |camel|
-              snake = camel_to_snake(camel)
-              map[snake] = camel
-              map[camel] = camel
+            {}.tap do |map|
+              camel_keys.each do |camel|
+                snake = camel_to_snake(camel)
+                map[snake] = camel
+                map[camel] = camel
+              end
             end
-            map
           end
 
           def camel_to_snake(str)

data/lib/ag_ui/run_input.rb

@@ -42,7 +42,7 @@ module AgUi
           validation_error = Protocol::JsonSchema::ValidationError.new(
             errors,
             definition_name: "RunAgentInput",
-            data: raw,
+            data:            raw,
           )
           raise InvalidError, validation_error.message
         end