ag-ui 0.2.0 → 1.0.0

data/lib/ag_ui/event_bridge.rb

@@ -6,42 +6,58 @@ require "ag_ui"
 module AgUi
   # The seam between brute's turn pipeline and the AG-UI SSE stream.
   #
-  # Brute middleware and terminal procs push `{type:, data:}` events into
-  # env[:events]; an EventBridge is that sink, translating each event into
-  # the matching typed SSE emitter as it arrives — the browser sees deltas
-  # mid-turn, not after.
+  # Brute middleware and terminal procs emit AG-UI events through the env; a
+  # bridge subscribes to that vocabulary on the agent's hooks and translates
+  # each one into the matching typed SSE emitter as it arrives — the browser
+  # sees deltas mid-turn, not after.
   #
-  #   pipeline.start(messages, events: AgUi::EventBridge.new(stream))
+  #   agent = Brute.agent.use(...).run(terminal)
+  #   AgUi::EventBridge.new(stream).subscribe(agent)
+  #   agent.start(messages)
   #
-  # Unknown event types (brute's own :log, :tool_result telemetry, etc.)
-  # are ignored — only the AG-UI vocabulary reaches the wire.
+  # Brute's own lifecycle events (:llm_start, :tool_end, the middleware and
+  # trace sets) are simply never subscribed to here, so only the AG-UI
+  # vocabulary reaches the wire.
   class EventBridge
     TRANSLATIONS = {
-      text_message_start: :translate_text_start,
-      text_message_content: :translate_text_content,
-      text_message_end: :translate_text_end,
-      tool_call_start: :translate_tool_call_start,
-      tool_call_args: :translate_tool_call_args,
-      tool_call_end: :translate_tool_call_end,
-      tool_call_result: :translate_tool_call_result,
-      activity_snapshot: :translate_activity_snapshot,
-      reasoning_start: :translate_reasoning_start,
-      reasoning_message_start: :translate_reasoning_message_start,
+      text_message_start:        :translate_text_start,
+      text_message_content:      :translate_text_content,
+      text_message_end:          :translate_text_end,
+      tool_call_start:           :translate_tool_call_start,
+      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,
+      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)
       @stream = stream
     end
 
-    def <<(event)
-      handler = TRANSLATIONS[event[:type]]
-      if handler
-        send(handler, event[:data] || {})
+    # Registers the whole vocabulary above on a hooks registry — or on
+    # anything else answering `on`, which a Brute pipeline builder does.
+    # Answers what it was given, so it composes in one line.
+    #
+    # Every subscriber takes (env, data, trace): brute hands the env first and
+    # appends the trace last, and a translation wants neither — the data is
+    # the whole of an AG-UI event.
+    def subscribe(hooks)
+      TRANSLATIONS.each do |event, handler|
+        hooks.on(event) { |_env, data, _trace| send(handler, data || {}) }
       end
-      self
+      hooks
     end
 
     private
@@ -64,8 +80,8 @@ module AgUi
 
       def translate_tool_call_start(data)
         @stream.tool_call_start(
-          tool_call_id: data[:tool_call_id],
-          tool_call_name: data[:tool_call_name],
+          tool_call_id:      data[:tool_call_id],
+          tool_call_name:    data[:tool_call_name],
           parent_message_id: data[:parent_message_id],
         )
       end
@@ -82,18 +98,52 @@ module AgUi
 
       def translate_tool_call_result(data)
         @stream.tool_call_result(
-          message_id: data[:message_id],
+          message_id:   data[:message_id],
           tool_call_id: data[:tool_call_id],
-          content: data[:content],
+          content:      data[:content],
         )
       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],
+          message_id:    data[:message_id],
           activity_type: data[:activity_type],
-          content: data[:content],
-          replace: data.fetch(:replace, true),
+          content:       data[:content],
+          replace:       data.fetch(:replace, true),
         )
       end
 
@@ -124,6 +174,15 @@ end
 __END__
 
 describe "AgUi::EventBridge" do
+  # A bridge is driven by emitting on a registry it has subscribed to, so the
+  # tests below build the pair and push events through the same path a turn
+  # would.
+  def bridged(stream)
+    hooks = Brute::Hooks::Registry.new
+    AgUi::EventBridge.new(stream).subscribe(hooks)
+    Brute::Hooks::Trace.new({}, hooks: hooks)
+  end
+
   read_frames = ->(stream) do
     frames = []
     while (chunk = stream.read)
@@ -134,12 +193,12 @@ describe "AgUi::EventBridge" do
 
   it "translates text events into SSE frames as they arrive" do
     stream = AgUi::Server::SSE::Stream.new(thread_id: "t1", run_id: "r1")
-    bridge = AgUi::EventBridge.new(stream)
+    env = bridged(stream)
 
-    bridge << { type: :text_message_start, data: { message_id: "m1" } }
-    bridge << { type: :text_message_content, data: { message_id: "m1", delta: "Hel" } }
-    bridge << { type: :text_message_content, data: { message_id: "m1", delta: "lo" } }
-    bridge << { type: :text_message_end, data: { message_id: "m1" } }
+    env.emit(:text_message_start, { message_id: "m1" })
+    env.emit(:text_message_content, { message_id: "m1", delta: "Hel" })
+    env.emit(:text_message_content, { message_id: "m1", delta: "lo" })
+    env.emit(:text_message_end, { message_id: "m1" })
     stream.finish
 
     frames = read_frames.(stream)
@@ -151,24 +210,65 @@ describe "AgUi::EventBridge" do
 
   it "drops empty deltas (protocol rule)" do
     stream = AgUi::Server::SSE::Stream.new(thread_id: "t1", run_id: "r1")
-    bridge = AgUi::EventBridge.new(stream)
+    env = bridged(stream)
 
-    bridge << { type: :text_message_content, data: { message_id: "m1", delta: "" } }
-    bridge << { type: :text_message_content, data: { message_id: "m1", delta: nil } }
+    env.emit(:text_message_content, { message_id: "m1", delta: "" })
+    env.emit(:text_message_content, { message_id: "m1", delta: nil })
     stream.finish
 
     read_frames.(stream).should == []
   end
 
-  it "ignores brute telemetry and unknown event types, returning self" do
+  it "leaves brute's own lifecycle events and unknown types off the wire" do
     stream = AgUi::Server::SSE::Stream.new(thread_id: "t1", run_id: "r1")
-    bridge = AgUi::EventBridge.new(stream)
+    env = bridged(stream)
 
-    result = bridge << { type: :log, data: { note: "internal" } }
-    bridge << { type: :whatever }
+    env.emit(Brute::Hooks::LLM_START_EVENT)
+    env.emit(:whatever, { note: "internal" })
     stream.finish
 
-    result.should.equal?(bridge)
     read_frames.(stream).should == []
   end
+
+  it "answers what it subscribed, so it composes in one line" do
+    hooks = Brute::Hooks::Registry.new
+    stream = AgUi::Server::SSE::Stream.new(thread_id: "t1", run_id: "r1")
+
+    AgUi::EventBridge.new(stream).subscribe(hooks).should.equal?(hooks)
+  end
+
+  it "translates shared-state events (snapshot + JSON-Patch delta)" do
+    stream = AgUi::Server::SSE::Stream.new(thread_id: "t1", run_id: "r1")
+    env = bridged(stream)
+
+    env.emit(:state_snapshot, { snapshot: { "theme" => "dark" } })
+    env.emit(
+      :state_delta,
+      { 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")
+    env = bridged(stream)
+
+    env.emit(:step_started, { step_name: "plan" })
+    env.emit(
+      :custom,
+      { name: "PredictState", value: [{ "state_key" => "document", "tool" => "write" }] },
+    )
+    env.emit(:step_finished, { 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/messages.rb

@@ -12,26 +12,24 @@ module AgUi
   module Messages
     class << self
       def to_brute(wire_messages)
-        log = Brute.log
-
-        wire_messages.each do |message|
-          case message["role"]
-          in "user"
-            log.user(text_content(message["content"]))
-          in "assistant"
-            log << assistant_message(message)
-          in "tool"
-            log.tool(message["content"].to_s, tool_call_id: message["toolCallId"])
-          in "system" | "developer"
-            log.system(message["content"].to_s)
-          in "activity" | "reasoning"
-            # Not part of the LLM conversation: activities render client-side
-            # (phase 4) and reasoning is provider-managed (phase 5).
-            nil
+        Brute.log.tap do |log|
+          wire_messages.each do |message|
+            case message["role"]
+            in "user"
+              log.user(text_content(message["content"]))
+            in "assistant"
+              log << assistant_message(message)
+            in "tool"
+              log.tool(message["content"].to_s, tool_call_id: message["toolCallId"])
+            in "system" | "developer"
+              log.system(message["content"].to_s)
+            in "activity" | "reasoning"
+              # Not part of the LLM conversation: activities render client-side
+              # (phase 4) and reasoning is provider-managed (phase 5).
+              nil
+            end
           end
         end
-
-        log
       end
 
       private
@@ -39,15 +37,15 @@ module AgUi
         def assistant_message(message)
           tool_calls = message["toolCalls"]&.map do |tc|
             Brute::ToolCall.new(
-              id: tc["id"],
-              name: tc.dig("function", "name"),
+              id:        tc["id"],
+              name:      tc.dig("function", "name"),
               arguments: parse_arguments(tc.dig("function", "arguments")),
             )
           end
 
           Brute::Message.new(
-            role: :assistant,
-            content: message["content"],
+            role:       :assistant,
+            content:    message["content"],
             tool_calls: tool_calls,
           )
         end

data/lib/ag_ui/middleware/a2ui.rb

@@ -29,37 +29,37 @@ module AgUi
     #
     # Sits OUTSIDE ToolRouter in the pipeline:
     #   use SystemPrompt; use A2ui, catalog: catalog; use ToolRouter
-    class A2ui
+    class A2ui < Brute::Middleware::Base
       TOOL_NAME = "render_a2ui"
 
       TOOL_DEFINITION = {
-        "name" => TOOL_NAME,
+        "name"        => TOOL_NAME,
         "description" =>
-          "Render a dynamic A2UI v0.9 surface with structured parameters. " \
+                         "Render a dynamic A2UI v0.9 surface with structured parameters. " \
           "Follow the A2UI render tool usage guide provided in context.",
-        "parameters" => {
-          "type" => "object",
+        "parameters"  => {
+          "type"       => "object",
           "properties" => {
-            "surfaceId" => {
-              "type" => "string",
+            "surfaceId"  => {
+              "type"        => "string",
               "description" => "Unique surface identifier.",
             },
             "components" => {
-              "type" => "array",
+              "type"        => "array",
               "description" =>
-                "A2UI v0.9 component array (flat format). The root component " \
+                               "A2UI v0.9 component array (flat format). The root component " \
                 "must have id \"root\".",
-              "items" => { "type" => "object" },
+              "items"       => { "type" => "object" },
             },
-            "data" => {
-              "type" => "object",
+            "data"       => {
+              "type"        => "object",
               "description" =>
-                "Initial data model for the surface. Written to the root path. " \
+                               "Initial data model for the surface. Written to the root path. " \
                 "Use for pre-filling form values (e.g. {\"form\": {\"name\": \"Alice\"}}) " \
                 "or providing data for components bound to data model paths.",
             },
           },
-          "required" => %w[surfaceId components],
+          "required"   => %w[surfaceId components],
         },
       }.freeze
 
@@ -123,7 +123,7 @@ module AgUi
           if catalog_components? && !metadata(env)[:a2ui_schema_injected]
             env[:messages].unshift(
               Brute::Message.new(
-                role: :system,
+                role:    :system,
                 content: "#{SCHEMA_CONTEXT_PREAMBLE}\n#{JSON.generate(@catalog.components)}",
               ),
             )
@@ -151,9 +151,9 @@ module AgUi
           surface_id = args["surfaceId"].to_s
 
           validation = ::AgUi::A2ui.validate_components(
-            components: args["components"],
-            data: args["data"].is_a?(Hash) ? args["data"] : {},
-            catalog: catalog_components? ? { "components" => @catalog.components } : nil,
+            components:        args["components"],
+            data:              args["data"].is_a?(Hash) ? args["data"] : {},
+            catalog:           catalog_components? ? { "components" => @catalog.components } : nil,
             validate_bindings: false,
           )
 
@@ -161,8 +161,8 @@ module AgUi
             render_failure(env, tool_call, validation["errors"])
           else
             ops = operations(args, surface_id, emitted_surfaces(env))
-            env[:events] << activity(tool_call, { "a2ui_operations" => ops })
-            env[:events] << result(tool_call, { "status" => "rendered" })
+            emit_activity(env, tool_call, { "a2ui_operations" => ops })
+            emit_result(env, tool_call, { "status" => "rendered" })
           end
         end
 
@@ -171,18 +171,19 @@ module AgUi
           metadata(env)[:a2ui_attempts] = attempt
           final = attempt >= MAX_ATTEMPTS
 
-          env[:events] << activity(
+          emit_activity(
+            env,
             tool_call,
             {
-              "status" => final ? "failed" : "retrying",
-              "attempt" => attempt,
+              "status"      => final ? "failed" : "retrying",
+              "attempt"     => attempt,
               "maxAttempts" => MAX_ATTEMPTS,
-              "errors" => errors,
+              "errors"      => errors,
             },
           )
 
           content = { "status" => "failed", "errors" => errors }
-          env[:events] << result(tool_call, content)
+          emit_result(env, tool_call, content)
           env[:messages].tool(JSON.generate(content), tool_call_id: tool_call.id)
 
           unless final
@@ -197,55 +198,55 @@ module AgUi
         # createSurface once per surface within the turn; updateComponents
         # always; updateDataModel when initial data was given.
         def operations(args, surface_id, emitted_surfaces)
-          ops = []
+          [].tap do |ops|
+            if emitted_surfaces.add?(surface_id)
+              ops << {
+                "version"       => "v0.9",
+                "createSurface" => { "surfaceId" => surface_id, "catalogId" => catalog_id },
+              }
+            end
 
-          if emitted_surfaces.add?(surface_id)
             ops << {
-              "version" => "v0.9",
-              "createSurface" => { "surfaceId" => surface_id, "catalogId" => catalog_id },
+              "version"          => "v0.9",
+              "updateComponents" => { "surfaceId" => surface_id, "components" => args["components"] },
             }
-          end
-
-          ops << {
-            "version" => "v0.9",
-            "updateComponents" => { "surfaceId" => surface_id, "components" => args["components"] },
-          }
 
-          if args["data"].is_a?(Hash) && !args["data"].empty?
-            ops << {
-              "version" => "v0.9",
-              "updateDataModel" => { "surfaceId" => surface_id, "path" => "/", "value" => args["data"] },
-            }
+            if args["data"].is_a?(Hash) && !args["data"].empty?
+              ops << {
+                "version"         => "v0.9",
+                "updateDataModel" => { "surfaceId" => surface_id, "path" => "/", "value" => args["data"] },
+              }
+            end
           end
-
-          ops
         end
 
         def catalog_id
           @default_catalog_id || BASIC_CATALOG_ID
         end
 
-        def activity(tool_call, content)
-          {
-            type: :activity_snapshot,
-            data: {
-              message_id: "a2ui-surface-#{tool_call.id}",
+        # One activity per tool call id, replacing rather than appending, so a
+        # surface re-rendered mid-turn updates in place on the client.
+        def emit_activity(env, tool_call, content)
+          env.emit(
+            :activity_snapshot,
+            {
+              message_id:    "a2ui-surface-#{tool_call.id}",
               activity_type: "a2ui-surface",
-              content: content,
-              replace: true,
+              content:       content,
+              replace:       true,
             },
-          }
+          )
         end
 
-        def result(tool_call, content)
-          {
-            type: :tool_call_result,
-            data: {
-              message_id: SecureRandom.uuid,
+        def emit_result(env, tool_call, content)
+          env.emit(
+            :tool_call_result,
+            {
+              message_id:   SecureRandom.uuid,
               tool_call_id: tool_call.id,
-              content: JSON.generate(content),
+              content:      JSON.generate(content),
             },
-          }
+          )
         end
     end
   end
@@ -254,6 +255,20 @@ end
 __END__
 
 describe "AgUi::Middleware::A2ui" do
+
+  # Brute 6 replaced the 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 A2ui speaks, in emission
+  # order, and accumulates across calls the way the old array did.
+  A2UI_EVENTS = %i[activity_snapshot tool_call_result].freeze
+
+  def recording(env, into)
+    hooks = Brute::Hooks::Registry.new
+    A2UI_EVENTS.each do |event|
+      hooks.on(event) { |_env, data, _trace| into << { type: event, data: data } }
+    end
+    Brute::Hooks::Trace.new(env, hooks: hooks)
+  end
   catalog = AgUi::A2ui::Catalog.new(
     catalog_id: "host://ai-catalog",
     components: { "Card" => { "description" => "A card", "props" => {} } },
@@ -292,12 +307,14 @@ describe "AgUi::Middleware::A2ui" do
       )
     end
 
-    env = { messages: Brute.log, events: [], tools: [] }
-    AgUi::Middleware::A2ui.new(terminal, catalog: id_only).call(env)
+    env = { messages: Brute.log, tools: [] }
+
+    events = []
+    AgUi::Middleware::A2ui.new(terminal, catalog: id_only).call(recording(env, events))
 
     seen[:messages].first.role.should == :assistant  # no schema system message
-    env[:events][0][:data][:content].key?("a2ui_operations").should == true
-    env[:events][0][:data][:content]["a2ui_operations"][0]["createSurface"]["catalogId"]
+    events[0][:data][:content].key?("a2ui_operations").should == true
+    events[0][:data][:content]["a2ui_operations"][0]["createSurface"]["catalogId"]
       .should == "app://cat"
   end
 
@@ -319,10 +336,12 @@ describe "AgUi::Middleware::A2ui" do
       )
     end
 
-    env = { messages: Brute.log, events: [], tools: [] }
-    AgUi::Middleware::A2ui.new(terminal, catalog: catalog).call(env)
+    env = { messages: Brute.log, tools: [] }
 
-    activity = env[:events][0]
+    events = []
+    AgUi::Middleware::A2ui.new(terminal, catalog: catalog).call(recording(env, events))
+
+    activity = events[0]
     activity[:type].should == :activity_snapshot
     activity[:data][:message_id].should == "a2ui-surface-tc1"
     activity[:data][:activity_type].should == "a2ui-surface"
@@ -333,7 +352,7 @@ describe "AgUi::Middleware::A2ui" do
     ops[1]["updateComponents"]["components"].first["id"].should == "root"
     ops[2]["updateDataModel"].should == { "surfaceId" => "s1", "path" => "/", "value" => { "title" => "hi" } }
 
-    result = env[:events][1]
+    result = events[1]
     result[:type].should == :tool_call_result
     result[:data][:tool_call_id].should == "tc1"
     result[:data][:content].should == "{\"status\":\"rendered\"}"
@@ -346,10 +365,12 @@ describe "AgUi::Middleware::A2ui" do
       env[:messages] << render_call.("surfaceId" => "s1", "components" => minimal_components)
     end
 
-    env = { messages: Brute.log, events: [], tools: [] }
-    AgUi::Middleware::A2ui.new(terminal).call(env)
+    env = { messages: Brute.log, tools: [] }
+
+    events = []
+    AgUi::Middleware::A2ui.new(terminal).call(recording(env, events))
 
-    ops = env[:events][0][:data][:content]["a2ui_operations"]
+    ops = events[0][:data][:content]["a2ui_operations"]
     ops[0]["createSurface"]["catalogId"].should ==
       "https://a2ui.org/specification/v0_9/basic_catalog.json"
   end
@@ -367,10 +388,12 @@ describe "AgUi::Middleware::A2ui" do
       )
     end
 
-    env = { messages: Brute.log, events: [], tools: [] }
-    AgUi::Middleware::A2ui.new(terminal, catalog: catalog).call(env)
+    env = { messages: Brute.log, tools: [] }
+
+    events = []
+    AgUi::Middleware::A2ui.new(terminal, catalog: catalog).call(recording(env, events))
 
-    activities = env[:events].select { |e| e[:type] == :activity_snapshot }
+    activities = events.select { |e| e[:type] == :activity_snapshot }
     first_ops = activities[0][:data][:content]["a2ui_operations"]
     second_ops = activities[1][:data][:content]["a2ui_operations"]
 
@@ -389,12 +412,14 @@ describe "AgUi::Middleware::A2ui" do
       )
     end
 
-    env = { messages: Brute.log, events: [], tools: [], should_exit: true }
-    AgUi::Middleware::A2ui.new(terminal, catalog: catalog).call(env)
+    env = { messages: Brute.log, tools: [], should_exit: true }
 
-    env[:events].length.should == 2
-    env[:events][0][:data][:content]["status"].should == "retrying"
-    JSON.parse(env[:events][1][:data][:content])["status"].should == "failed"
+    events = []
+    AgUi::Middleware::A2ui.new(terminal, catalog: catalog).call(recording(env, events))
+
+    events.length.should == 2
+    events[0][:data][:content]["status"].should == "retrying"
+    JSON.parse(events[1][:data][:content])["status"].should == "failed"
     env[:messages].last.role.should == :tool
     env[:messages].last.tool_call_id.should == "tc2"
     env[:should_exit].should == false
@@ -417,11 +442,13 @@ describe "AgUi::Middleware::A2ui" do
       env[:messages] << invalid.("tc#{calls}")
     end
 
-    env = { messages: Brute.log, events: [], tools: [], metadata: {} }
+    env = { messages: Brute.log, tools: [], metadata: {} }
+
+    events = []
     mw = AgUi::Middleware::A2ui.new(terminal, catalog: catalog)
 
-    mw.call(env)
-    first = env[:events][0][:data][:content]
+    mw.call(recording(env, events))
+    first = events[0][:data][:content]
     first["status"].should == "retrying"
     first["attempt"].should == 1
     codes = first["errors"].map { |e| e["code"] }
@@ -430,11 +457,11 @@ describe "AgUi::Middleware::A2ui" do
     env[:should_exit].should == false
 
     env[:should_exit] = true
-    mw.call(env)
+    mw.call(recording(env, events))
     env[:should_exit] = true
-    mw.call(env)
+    mw.call(recording(env, events))
 
-    statuses = env[:events].select { |e| e[:type] == :activity_snapshot }
+    statuses = events.select { |e| e[:type] == :activity_snapshot }
       .map { |e| e[:data][:content]["status"] }
     statuses.should == %w[retrying retrying failed]
     env[:should_exit].should == true