ag-ui 0.1.0

data/lib/ag_ui/protocol/json_schema.rb

@@ -0,0 +1,270 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "ag_ui"
+
+module AgUi
+  module Protocol
+    # Schema-validated AG-UI protocol objects, powered by json_schemer.
+    # Same mechanism as a2a's Protocol::JsonSchema, over the schema bundle
+    # generated from the reference Python SDK's pydantic models
+    # (data/generate-ag-ui-schema.py -> data/ag_ui.json).
+    #
+    #   AgUi::Protocol::JsonSchema["TextMessageContentEvent"]
+    #   #=> Class < Definition with .schema, .valid?, reader methods
+    #
+    #   AgUi::Protocol::JsonSchema.list_definitions
+    #   #=> ["ActivityDeltaEvent", "ActivityMessage", ...]
+    #
+    module JsonSchema
+      DATA_PATH = File.expand_path("../../../data/ag_ui.json", __dir__).freeze
+
+      @definition_classes = {}
+      @schemer = nil
+      @raw_schema = nil
+
+      class << self
+        # Look up a definition by model class name.
+        #
+        #   AgUi::Protocol::JsonSchema["RunAgentInput"]
+        #   #=> Class < Definition
+        #
+        def [](name)
+          @definition_classes[name] ||= begin
+            definitions = raw_schema.fetch("definitions", {})
+
+            unless definitions.key?(name)
+              raise "No AG-UI definition found for #{name.inspect}!" \
+                "\nAvailable: #{list_definitions.join(", ")}"
+            end
+
+            ref_schema = schemer.ref("#/definitions/#{name}")
+            build_definition_class(ref_schema, name, definitions[name])
+          end
+        end
+
+        # All available definition names, sorted.
+        def list_definitions
+          raw_schema.fetch("definitions", {}).keys.sort
+        end
+
+        # Definition names for concrete events (have a const "type"),
+        # mapped from wire event type to definition name:
+        #
+        #   { "TEXT_MESSAGE_CONTENT" => "TextMessageContentEvent", ... }
+        #
+        def event_types
+          @event_types ||= raw_schema.fetch("definitions", {}).each_with_object({}) do |(name, defn), map|
+            const = defn.dig("properties", "type", "const")
+            if const
+              map[const] = name
+            end
+          end.freeze
+        end
+
+        # The JSONSchemer instance for the full AG-UI schema bundle.
+        def schemer
+          @schemer ||= JSONSchemer.schema(raw_schema)
+        end
+
+        # The parsed schema hash (all $refs are already internal —
+        # the generator emits #/definitions/... pointers).
+        def raw_schema
+          @raw_schema ||= JSON.parse(File.read(DATA_PATH))
+        end
+
+        # Reset all cached state (useful for tests).
+        def reset!
+          @definition_classes.clear
+          @schemer = nil
+          @raw_schema = nil
+          @event_types = nil
+        end
+
+        private
+
+          # Build a Definition subclass for a specific AG-UI type.
+          def build_definition_class(schema_instance, definition_name, raw_definition)
+            properties     = raw_definition.fetch("properties", {})
+            camel_keys     = properties.keys
+            snake_to_camel = build_snake_to_camel(camel_keys)
+            prop_refs      = build_property_refs(properties)
+
+            reader_pairs = camel_keys.map { |ck| [camel_to_snake(ck).to_sym, ck] }
+
+            Class.new(Definition) do
+              @schema            = schema_instance
+              @definition_name   = definition_name
+              @schema_properties = camel_keys
+              @snake_to_camel    = snake_to_camel
+              @property_refs     = prop_refs
+
+              class << self
+                def schema           = @schema
+                def definition_name  = @definition_name
+                def schema_properties = @schema_properties
+                def snake_to_camel_map = @snake_to_camel
+                def property_refs    = @property_refs
+              end
+
+              reader_pairs.each do |snake_sym, camel_key|
+                define_method(snake_sym) { @data[camel_key] }
+              end
+            end
+          end
+
+          # Inspect each property's schema for $ref pointers and build a map
+          # of { camelKey => [:kind, "DefinitionName"] } so Definition can
+          # auto-wrap nested Hashes. Pydantic wraps optional properties in
+          # anyOf [<schema>, {type: null}], so detection looks through the
+          # first non-null anyOf variant.
+          #
+          #   :object — direct $ref
+          #   :array  — items.$ref
+          #   :map    — additionalProperties.$ref
+          def build_property_refs(properties)
+            definitions = raw_schema.fetch("definitions", {})
+            refs = {}
+
+            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
+
+              if ref
+                name = ref_name_for(ref)
+                if name && definitions.dig(name, "properties")
+                  refs[camel_key] = [kind, name]
+                end
+              end
+            end
+
+            refs
+          end
+
+          # Pydantic optional: anyOf [<schema>, {"type": "null"}] — return
+          # the non-null variant (only when the anyOf is that exact shape).
+          def unwrap_optional(prop_schema)
+            variants = prop_schema["anyOf"]
+            if variants
+              non_null = variants.reject { |s| s == { "type" => "null" } }
+              non_null.length == 1 ? non_null.first : prop_schema
+            else
+              prop_schema
+            end
+          end
+
+          def ref_name_for(ref)
+            if ref.start_with?("#/definitions/")
+              ref.sub("#/definitions/", "")
+            end
+          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
+            end
+            map
+          end
+
+          def camel_to_snake(str)
+            str.gsub(/([A-Z])/) { "_#{$1.downcase}" }
+               .delete_prefix("_")
+          end
+      end
+    end
+  end
+end
+
+__END__
+
+describe "AgUi::Protocol::JsonSchema" do
+  schema = AgUi::Protocol::JsonSchema
+
+  it "loads the raw schema with definitions" do
+    schema.raw_schema["definitions"].should.be.kind_of(Hash)
+  end
+
+  it "contains only internal $refs" do
+    external = []
+    walk = ->(obj) do
+      case obj
+      when Hash
+        obj.each do |k, v|
+          if k == "$ref" && v.is_a?(String) && !v.start_with?("#")
+            external << v
+          else
+            walk.(v)
+          end
+        end
+      when Array
+        obj.each { |v| walk.(v) }
+      end
+    end
+    walk.(schema.raw_schema["definitions"])
+    external.should == []
+  end
+
+  it "lists definitions sorted, covering the core vocabulary" do
+    defs = schema.list_definitions
+    defs.should == defs.sort
+    %w[RunAgentInput RunStartedEvent RunFinishedEvent TextMessageContentEvent
+       ToolCallStartEvent ActivitySnapshotEvent AssistantMessage ToolMessage
+       AgentCapabilities].each { |d| defs.should.include?(d) }
+  end
+
+  it "maps wire event types to definition names" do
+    schema.event_types["TEXT_MESSAGE_CONTENT"].should == "TextMessageContentEvent"
+    schema.event_types["RUN_STARTED"].should == "RunStartedEvent"
+    schema.event_types["ACTIVITY_SNAPSHOT"].should == "ActivitySnapshotEvent"
+  end
+
+  it "returns cached Definition subclasses" do
+    a = schema["RunStartedEvent"]
+    b = schema["RunStartedEvent"]
+    (a < AgUi::Protocol::JsonSchema::Definition).should == true
+    a.object_id.should == b.object_id
+  end
+
+  it "builds definitions that serialize snake_case input to camelCase wire form" do
+    event = schema["TextMessageContentEvent"].new(
+      type: "TEXT_MESSAGE_CONTENT", message_id: "m1", delta: "Hi"
+    )
+    event.valid?.should == true
+    event.to_h.should == { "type" => "TEXT_MESSAGE_CONTENT", "messageId" => "m1", "delta" => "Hi" }
+  end
+
+  it "rejects invalid events" do
+    event = schema["TextMessageContentEvent"].new(type: "TEXT_MESSAGE_CONTENT", message_id: "m1")
+    event.valid?.should == false
+  end
+
+  it "detects pydantic-optional anyOf refs as :object property refs" do
+    schema["RunStartedEvent"].property_refs["input"].should == [:object, "RunAgentInput"]
+  end
+
+  it "validates RunAgentInput built via Definition" do
+    input = schema["RunAgentInput"].new(
+      thread_id: "t1", run_id: "r1", state: {}, messages: [],
+      tools: [], context: [], forwarded_props: {}
+    )
+    input.valid?.should == true
+    input.to_h["threadId"].should == "t1"
+  end
+
+  it "validates a raw wire hash with explicit nulls via schemer (inbound path)" do
+    raw = { "threadId" => "t1", "runId" => "r1", "state" => nil, "messages" => [],
+            "tools" => [], "context" => [], "forwardedProps" => nil, "extra" => 1 }
+    schema.schemer.ref("#/definitions/RunAgentInput").valid?(raw).should == true
+  end
+end

data/lib/ag_ui/run_input.rb

@@ -0,0 +1,149 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "ag_ui"
+
+module AgUi
+  # Parses and validates the POST body of /agent/:id/run — the
+  # RunAgentInput envelope (threadId, runId, messages, tools, context,
+  # forwardedProps, state, resume?).
+  #
+  # Validation runs against the RAW parsed hash so explicit nulls
+  # (state: null, forwardedProps: null) survive — Definition#to_h
+  # compacts nils away, which would fail required-field checks.
+  # The returned Definition gives snake_case readers and pattern
+  # matching over the same data.
+  #
+  #   input = AgUi::RunInput.parse(request_body)
+  #   input.thread_id  #=> "t1"
+  #   input.messages   #=> [Definition(UserMessage), ...]
+  #
+  module RunInput
+    # Raised on malformed JSON or schema-invalid input. The route layer
+    # renders this as 400 {"error": "Invalid request body", "details": ...}.
+    class InvalidError < StandardError; end
+
+    SCHEMA_REF = "#/definitions/RunAgentInput"
+
+    class << self
+      def parse(body)
+        begin
+          raw = JSON.parse(body)
+        rescue JSON::ParserError => e
+          raise InvalidError, "malformed JSON: #{e.message}"
+        end
+
+        unless raw.is_a?(Hash)
+          raise InvalidError, "expected a JSON object, got #{raw.class}"
+        end
+
+        errors = schema.validate(raw).to_a
+        unless errors.empty?
+          validation_error = Protocol::JsonSchema::ValidationError.new(
+            errors,
+            definition_name: "RunAgentInput",
+            data: raw,
+          )
+          raise InvalidError, validation_error.message
+        end
+
+        Protocol::JsonSchema["RunAgentInput"].new(raw)
+      end
+
+      def schema
+        @schema ||= Protocol::JsonSchema.schemer.ref(SCHEMA_REF)
+      end
+    end
+  end
+end
+
+__END__
+
+describe "AgUi::RunInput" do
+  minimal = {
+    "threadId" => "t1", "runId" => "r1", "state" => nil,
+    "messages" => [], "tools" => [], "context" => [], "forwardedProps" => nil,
+  }
+
+  it "parses a minimal valid body with explicit nulls" do
+    input = AgUi::RunInput.parse(JSON.generate(minimal))
+    input.thread_id.should == "t1"
+    input.run_id.should == "r1"
+    input.messages.should == []
+  end
+
+  it "parses the full envelope: messages, tools, context" do
+    body = minimal.merge(
+      "messages" => [
+        { "id" => "u1", "role" => "user", "content" => "hi" },
+        { "id" => "a1", "role" => "assistant",
+          "toolCalls" => [{ "id" => "tc1", "type" => "function",
+                            "function" => { "name" => "navigate", "arguments" => "{}" } }] },
+        { "id" => "tr1", "role" => "tool", "content" => "{}", "toolCallId" => "tc1" },
+      ],
+      "tools" => [{ "name" => "navigate", "description" => "Go to a page",
+                    "parameters" => { "type" => "object" } }],
+      "context" => [{ "description" => "currentPath", "value" => "/data" }],
+    )
+
+    input = AgUi::RunInput.parse(JSON.generate(body))
+    input.messages.length.should == 3
+    input.tools.first.name.should == "navigate"
+    input.context.first.value.should == "/data"
+
+    # Messages are a discriminated union — they stay raw camelCase hashes.
+    input.messages.last["toolCallId"].should == "tc1"
+
+    case input
+    in { thread_id: String => tid }
+      tid.should == "t1"
+    end
+  end
+
+  it "accepts multimodal user message content parts" do
+    body = minimal.merge(
+      "messages" => [
+        { "id" => "u1", "role" => "user", "content" => [
+          { "type" => "text", "text" => "what is this?" },
+          { "type" => "image",
+            "source" => { "type" => "data", "value" => "aGk=", "mimeType" => "image/png" } },
+        ] },
+      ],
+    )
+
+    input = AgUi::RunInput.parse(JSON.generate(body))
+    input.messages.first["content"].length.should == 2
+  end
+
+  it "tolerates unknown extra fields (extra=allow)" do
+    input = AgUi::RunInput.parse(JSON.generate(minimal.merge("lastSeenEventId" => "5")))
+    input.thread_id.should == "t1"
+  end
+
+  it "rejects malformed JSON" do
+    begin
+      AgUi::RunInput.parse("{nope")
+      raise "expected InvalidError"
+    rescue AgUi::RunInput::InvalidError => e
+      e.message.should.include?("malformed JSON")
+    end
+  end
+
+  it "rejects non-object bodies" do
+    begin
+      AgUi::RunInput.parse("[1,2]")
+      raise "expected InvalidError"
+    rescue AgUi::RunInput::InvalidError => e
+      e.message.should.include?("expected a JSON object")
+    end
+  end
+
+  it "rejects schema-invalid bodies with readable details" do
+    begin
+      AgUi::RunInput.parse(JSON.generate({ "threadId" => "t1" }))
+      raise "expected InvalidError"
+    rescue AgUi::RunInput::InvalidError => e
+      e.message.should.include?("runId")
+    end
+  end
+end

data/lib/ag_ui/run_loop.rb

@@ -0,0 +1,285 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "console"
+require "ag_ui"
+
+module AgUi
+  # The run loop: drives one AG-UI run through a brute turn pipeline.
+  #
+  # The terminal block is the LLM call (ruby_llm or anything else) — it
+  # receives brute's env ({messages:, events:, ...}), streams deltas into
+  # env[:events] (translated live to SSE by EventBridge) and appends the
+  # assistant message to env[:messages]. The gem stays LLM-agnostic.
+  #
+  #   run_loop = AgUi::RunLoop.new(system_prompt: PROMPT) do |env|
+  #     env[:events] << { type: :text_message_start, data: { message_id: id } }
+  #     ... stream provider chunks ...
+  #     env[:events] << { type: :text_message_end, data: { message_id: id } }
+  #     env[:messages].assistant(full_text)
+  #   end
+  #
+  #   app = AgUi.agent(agent_id: "default", &run_loop)
+  #
+  # Lifecycle per run (the AG-UI contract): RUN_STARTED first; then the
+  # pipeline streams; then RUN_FINISHED — or RUN_ERROR if the pipeline
+  # raised. Schema-validation failures are wire-contract bugs and re-raise
+  # after RUN_ERROR so they fail loudly in dev.
+  class RunLoop
+    # a2ui: nil/false = off; an AgUi::A2ui::Catalog = on with that catalog;
+    # true = on degraded (tool injected, no component schema).
+    # server_tools: [{name:, description:, parameters:, handler:}] execute
+    # inline and the turn loops (Loop::ToolResult) until the model answers
+    # in text, a client tool defers to the browser, or max_iterations hits.
+    def initialize(system_prompt: nil, validate: true, middleware: [], a2ui: nil,
+                   server_tools: [], max_iterations: 10, &terminal)
+      unless terminal
+        raise ArgumentError, "RunLoop requires a terminal block (the LLM call)"
+      end
+
+      @system_prompt = system_prompt
+      @validate = validate
+      @middleware = middleware
+      @a2ui = a2ui
+      @server_tools = server_tools
+      @max_iterations = max_iterations
+      @terminal = terminal
+    end
+
+    def a2ui_enabled?
+      !(@a2ui.nil? || @a2ui == false)
+    end
+
+    # AgUi.agent takes a block; RunLoop quacks like one.
+    def to_proc
+      run_loop = self
+      proc { |rack_env| run_loop.handle(rack_env) }
+    end
+
+    def handle(rack_env)
+      input = rack_env["ag_ui.input"]
+
+      rack_env["ag_ui.stream"].open(
+        thread_id: input.thread_id,
+        run_id: input.run_id,
+        validate: @validate,
+      ) do |stream|
+        stream.run_started
+        run(stream, input)
+      end
+    end
+
+    private
+
+      def run(stream, input)
+        pipeline(input).start(
+          Messages.to_brute(input.messages),
+          events: EventBridge.new(stream),
+        )
+        stream.run_finished
+      rescue Protocol::JsonSchema::ValidationError => e
+        Console.error(self, "wire-contract bug: #{e.message}", e)
+        stream.run_error(message: "Internal error", code: "validation")
+        raise
+      rescue => e
+        Console.error(self, "run failed: #{e.class}: #{e.message}", e)
+        stream.run_error(message: e.message, code: e.class.name)
+      end
+
+      def pipeline(input)
+        agent = Brute.agent
+        agent.use(
+          AgUi::Middleware::SystemPrompt,
+          prompt: @system_prompt,
+          context: input.context,
+        )
+        agent.use(AgUi::Middleware::ForwardedProps, props: input.forwarded_props)
+        agent.use(Brute::Middleware::Loop::ToolResult)
+        agent.use(Brute::Middleware::MaxIterations, max_iterations: @max_iterations)
+        if a2ui_enabled?
+          catalog = @a2ui.is_a?(AgUi::A2ui::Catalog) ? @a2ui : nil
+          agent.use(AgUi::Middleware::A2ui, catalog: catalog)
+        end
+        agent.use(AgUi::Middleware::ToolRouter, tools: input.tools, server_tools: @server_tools)
+        @middleware.each do |(klass, options)|
+          agent.use(klass, **(options || {}))
+        end
+        agent.run(@terminal)
+      end
+  end
+end
+
+__END__
+
+describe "AgUi::RunLoop" do
+  minimal_input = JSON.generate({
+    "threadId" => "t1", "runId" => "r1", "state" => nil,
+    "messages" => [{ "id" => "u1", "role" => "user", "content" => "hi" }],
+    "tools" => [], "context" => [{ "description" => "currentPath", "value" => "/data" }],
+    "forwardedProps" => nil,
+  })
+
+  request = ->(app, body) do
+    app.call({
+      "REQUEST_METHOD" => "POST",
+      "PATH_INFO" => "/agent/default/run",
+      "rack.input" => StringIO.new(body),
+    })
+  end
+
+  read_frames = ->(stream) do
+    frames = []
+    while (chunk = stream.read)
+      frames << JSON.parse(chunk.sub(/\Adata: /, "").strip)
+    end
+    frames
+  end
+
+  it "streams a full run through a brute pipeline" do
+    seen_env = nil
+    run_loop = AgUi::RunLoop.new(system_prompt: "Be terse.") do |env|
+      seen_env = env
+      env[:events] << { type: :text_message_start, data: { message_id: "m1" } }
+      env[:events] << { type: :text_message_content, data: { message_id: "m1", delta: "Hello" } }
+      env[:events] << { type: :text_message_end, data: { message_id: "m1" } }
+      env[:messages].assistant("Hello")
+    end
+
+    app = AgUi.agent(agent_id: "default", &run_loop)
+    _status, _headers, body = request.(app, minimal_input)
+
+    frames = read_frames.(body)
+    frames.map { |f| f["type"] }.should == %w[
+      RUN_STARTED TEXT_MESSAGE_START TEXT_MESSAGE_CONTENT TEXT_MESSAGE_END RUN_FINISHED
+    ]
+
+    # The pipeline saw: system prompt (with context addendum) + history.
+    seen_env[:messages].map(&:role).should == [:system, :user, :assistant]
+    seen_env[:messages].first.content.should.include?("Be terse.")
+    seen_env[:messages].first.content.should.include?("currentPath: /data")
+  end
+
+  it "ends the run with RUN_ERROR when the terminal raises" do
+    run_loop = AgUi::RunLoop.new { |_env| raise "provider exploded" }
+    app = AgUi.agent(agent_id: "default", &run_loop)
+
+    _status, _headers, body = request.(app, minimal_input)
+    frames = read_frames.(body)
+
+    frames.map { |f| f["type"] }.should == %w[RUN_STARTED RUN_ERROR]
+    frames.last["message"].should == "provider exploded"
+    frames.last["code"].should == "RuntimeError"
+  end
+
+  it "supports extra brute middleware" do
+    marker = Class.new do
+      def initialize(app, note:)
+        @app = app
+        @note = note
+      end
+
+      def call(env)
+        env[:metadata][:note] = @note
+        @app.call(env)
+      end
+    end
+
+    seen = nil
+    run_loop = AgUi::RunLoop.new(middleware: [[marker, { note: "hi" }]]) do |env|
+      seen = env[:metadata][:note]
+      env[:messages].assistant("ok")
+    end
+
+    request.(AgUi.agent(&run_loop), minimal_input)
+    seen.should == "hi"
+  end
+
+  it "requires a terminal block" do
+    lambda { AgUi::RunLoop.new }.should.raise(ArgumentError)
+  end
+
+  it "executes server tools inline and loops the turn to completion" do
+    weather_tool = {
+      name: "get_weather",
+      description: "Weather for a city",
+      parameters: { "type" => "object" },
+      handler: ->(args) { { "city" => args["city"], "temp" => 21 } },
+    }
+
+    iterations = 0
+    terminal = ->(env) do
+      iterations += 1
+      if iterations == 1
+        env[:messages] << Brute::Message.new(
+          role: :assistant, content: nil,
+          tool_calls: [{ id: "tc1", name: "get_weather", arguments: { "city" => "Lisbon" } }],
+        )
+      else
+        # The model sees its own call + the executed result.
+        env[:messages].last.role.should == :tool
+        env[:messages].last.content.should == "{\"city\":\"Lisbon\",\"temp\":21}"
+        env[:events] << { type: :text_message_start, data: { message_id: "m2" } }
+        env[:events] << { type: :text_message_content, data: { message_id: "m2", delta: "21C in Lisbon" } }
+        env[:events] << { type: :text_message_end, data: { message_id: "m2" } }
+        env[:messages].assistant("21C in Lisbon")
+      end
+    end
+
+    run_loop = AgUi::RunLoop.new(server_tools: [weather_tool], &terminal)
+    app = AgUi.agent(agent_id: "default", &run_loop)
+
+    _status, _headers, stream = request.(app, minimal_input)
+    frames = read_frames.(stream)
+
+    iterations.should == 2
+    frames.map { |f| f["type"] }.should == %w[
+      RUN_STARTED
+      TOOL_CALL_START TOOL_CALL_ARGS TOOL_CALL_END TOOL_CALL_RESULT
+      TEXT_MESSAGE_START TEXT_MESSAGE_CONTENT TEXT_MESSAGE_END
+      RUN_FINISHED
+    ]
+    frames[4]["content"].should == "{\"city\":\"Lisbon\",\"temp\":21}"
+  end
+
+  it "streams the full A2UI sequence when the model renders a surface" do
+    catalog = AgUi::A2ui::Catalog.new(
+      catalog_id: "host://ai-catalog",
+      components: { "Card" => {} },
+    )
+
+    terminal = ->(env) do
+      env[:tools].map { |t| t["name"] }.should.include?("render_a2ui")
+      env[:messages].first.content.should.include?("A2UI Component Schema")
+
+      env[:messages] << Brute::Message.new(
+        role: :assistant, content: nil,
+        tool_calls: [{ id: "tc9", name: "render_a2ui", arguments: {
+          "surfaceId" => "s1",
+          "components" => [{ "id" => "root", "component" => "Card" }],
+        } }],
+      )
+    end
+
+    run_loop = AgUi::RunLoop.new(a2ui: catalog, &terminal)
+    app = AgUi.agent(agent_id: "default", &run_loop)
+
+    _status, _headers, stream = request.(app, minimal_input)
+    frames = read_frames.(stream)
+
+    frames.map { |f| f["type"] }.should == %w[
+      RUN_STARTED TOOL_CALL_START TOOL_CALL_ARGS TOOL_CALL_END
+      ACTIVITY_SNAPSHOT TOOL_CALL_RESULT RUN_FINISHED
+    ]
+
+    snapshot = frames[4]
+    snapshot["messageId"].should == "a2ui-surface-tc9"
+    snapshot["activityType"].should == "a2ui-surface"
+    snapshot["replace"].should == true
+    snapshot["content"]["a2ui_operations"][0]["createSurface"]["catalogId"]
+      .should == "host://ai-catalog"
+
+    result = frames[5]
+    result["toolCallId"].should == "tc9"
+    result["content"].should == "{\"status\":\"rendered\"}"
+  end
+end