agent2agent 1.0.8 → 1.1.0

data/lib/a2a/middleware/sse_stream.rb

@@ -0,0 +1,235 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "a2a"
+require "a2a/sse"
+require "async"
+
+module A2A
+  module Middleware
+    # Sets up an SSE stream builder on `env["a2a.stream"]`.
+    #
+    # The builder detects the protocol binding (REST vs JSON-RPC) from
+    # env and creates the correct stream subclass when `open` is called.
+    # The `open` block runs inside an Async fiber and the stream is
+    # automatically finished when the block exits (even on exception).
+    #
+    # If the handler never calls `open`, the builder is removed from env
+    # so the binding layer doesn't mistake it for a real stream.
+    #
+    # Usage:
+    #
+    #   on "SendStreamingMessage" do
+    #     use A2A::Middleware::SSEStream
+    #     use A2A::Middleware::ExtractMessage
+    #     respond_with -> (env) {
+    #       env["a2a.stream"].open(task_id: "t1", context_id: "c1") do |s|
+    #         s.task(status: { state: "TASK_STATE_WORKING" })
+    #         s.status_update(status: { state: "TASK_STATE_COMPLETED" })
+    #       end
+    #     }
+    #   end
+    #
+    class SSEStream
+      def initialize(app)
+        @app = app
+      end
+
+      def call(env)
+        builder = StreamBuilder.new(env)
+        env["a2a.stream"] = builder
+
+        result = @app.call(env)
+
+        # If open was never called, clear the builder so the binding
+        # layer doesn't mistake it for a real stream.
+        env.delete("a2a.stream") if env["a2a.stream"].equal?(builder)
+
+        result
+      end
+    end
+
+    # Factory that creates the correct SSE stream subclass based on the
+    # protocol binding, then runs the caller's block inside Async with
+    # automatic finish on exit.
+    #
+    # Created by SSEStream middleware — not intended for direct use.
+    #
+    class StreamBuilder
+      def initialize(env)
+        @env = env
+      end
+
+      # Create and open an SSE stream for the current request.
+      #
+      # Detects REST vs JSON-RPC from the env, constructs the correct
+      # stream subclass, and yields it to the block. The block runs
+      # inside an Async fiber. The stream is automatically finished
+      # when the block exits, even if an exception is raised.
+      #
+      # @param task_id    [String] the task identifier for this stream
+      # @param context_id [String] the context identifier for this stream
+      # @yieldparam stream [A2A::SSE::Stream] the opened stream
+      # @return [nil]
+      #
+      def open(task_id:, context_id:, &block)
+        stream = if @env["a2a.json_rpc_id"]
+          A2A::SSE::JsonRpcStream.new(
+            task_id: task_id, context_id: context_id,
+            json_rpc_id: @env["a2a.json_rpc_id"]
+          )
+        else
+          A2A::SSE::RestStream.new(
+            task_id: task_id, context_id: context_id
+          )
+        end
+
+        @env["a2a.stream"] = stream
+
+        Async do
+          block.call(stream)
+        ensure
+          stream.finish
+        end
+
+        nil
+      end
+    end
+  end
+end
+
+test do
+  describe "A2A::Middleware::SSEStream" do
+    it "sets a StreamBuilder on env[\"a2a.stream\"]" do
+      seen_stream = nil
+      downstream = ->(env) { seen_stream = env["a2a.stream"]; :ok }
+
+      mw = A2A::Middleware::SSEStream.new(downstream)
+      env = {}
+      mw.call(env)
+
+      seen_stream.should.be.kind_of(A2A::Middleware::StreamBuilder)
+    end
+
+    it "auto-clears builder if open was never called" do
+      downstream = ->(env) { :ok }
+
+      mw = A2A::Middleware::SSEStream.new(downstream)
+      env = {}
+      mw.call(env)
+
+      env.key?("a2a.stream").should == false
+    end
+
+    it "preserves env[\"a2a.stream\"] when open was called" do
+      downstream = ->(env) {
+        env["a2a.stream"].open(task_id: "t1", context_id: "c1") do |s|
+          # no-op
+        end
+      }
+
+      mw = A2A::Middleware::SSEStream.new(downstream)
+      env = {}
+      mw.call(env)
+
+      env["a2a.stream"].should.be.kind_of(A2A::SSE::RestStream)
+    end
+
+    it "returns the downstream result" do
+      downstream = ->(env) { :result_value }
+
+      mw = A2A::Middleware::SSEStream.new(downstream)
+      result = mw.call({})
+
+      result.should == :result_value
+    end
+  end
+
+  describe "A2A::Middleware::StreamBuilder" do
+    it "creates a RestStream when no JSON-RPC ID" do
+      env = {}
+      builder = A2A::Middleware::StreamBuilder.new(env)
+
+      builder.open(task_id: "t1", context_id: "c1") do |s|
+        s.should.be.kind_of(A2A::SSE::RestStream)
+      end
+
+      env["a2a.stream"].should.be.kind_of(A2A::SSE::RestStream)
+    end
+
+    it "creates a JsonRpcStream when JSON-RPC ID is present" do
+      env = { "a2a.json_rpc_id" => 42 }
+      builder = A2A::Middleware::StreamBuilder.new(env)
+
+      builder.open(task_id: "t1", context_id: "c1") do |s|
+        s.should.be.kind_of(A2A::SSE::JsonRpcStream)
+      end
+
+      env["a2a.stream"].should.be.kind_of(A2A::SSE::JsonRpcStream)
+    end
+
+    it "passes task_id and context_id to the stream" do
+      env = {}
+      builder = A2A::Middleware::StreamBuilder.new(env)
+
+      builder.open(task_id: "t1", context_id: "c1") do |s|
+        s.task_id.should == "t1"
+        s.context_id.should == "c1"
+      end
+    end
+
+    it "auto-finishes the stream when block completes" do
+      env = {}
+      builder = A2A::Middleware::StreamBuilder.new(env)
+
+      builder.open(task_id: "t1", context_id: "c1") do |s|
+        s.task(status: { state: "TASK_STATE_WORKING" })
+      end
+
+      stream = env["a2a.stream"]
+      # After finish, read drains remaining data then returns nil
+      stream.read  # the task event
+      stream.read.should.be.nil
+    end
+
+    it "auto-finishes even when block raises" do
+      env = {}
+      builder = A2A::Middleware::StreamBuilder.new(env)
+
+      builder.open(task_id: "t1", context_id: "c1") do |s|
+        s.task(status: { state: "TASK_STATE_WORKING" })
+        raise "boom"
+      end
+
+      stream = env["a2a.stream"]
+      stream.read  # the task event
+      stream.read.should.be.nil
+    end
+
+    it "returns nil" do
+      env = {}
+      builder = A2A::Middleware::StreamBuilder.new(env)
+
+      result = builder.open(task_id: "t1", context_id: "c1") do |s|
+        # no-op
+      end
+
+      result.should.be.nil
+    end
+
+    it "typed methods inject task_id and context_id" do
+      env = {}
+      builder = A2A::Middleware::StreamBuilder.new(env)
+
+      builder.open(task_id: "t1", context_id: "c1") do |s|
+        s.status_update(status: { state: "TASK_STATE_COMPLETED", timestamp: "2025-01-01T00:00:00Z" })
+      end
+
+      stream = env["a2a.stream"]
+      chunk = stream.read
+      parsed = JSON.parse(chunk.sub(/\Adata: /, "").strip)
+      parsed["statusUpdate"]["taskId"].should == "t1"
+      parsed["statusUpdate"]["contextId"].should == "c1"
+    end
+  end
+end

data/lib/a2a/middleware.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+require_relative "middleware/fetch_task"
+require_relative "middleware/limit_history_length"
+require_relative "middleware/limit_pagination_size"
+require_relative "middleware/extract_message"
+require_relative "middleware/sse_stream"

data/lib/a2a/schema/definition.rb

@@ -27,6 +27,7 @@ module A2A
       def initialize(hash = {})
         props    = self.class.schema_properties
         snake    = self.class.snake_to_camel_map
+        refs     = self.class.property_refs
         @data    = {}
 
         hash.each do |key, value|
@@ -36,7 +37,13 @@ module A2A
           camel = snake[k] || k
 
           if props.include?(camel)
-            @data[camel] = value.is_a?(Definition) ? value.to_h : value
+            @data[camel] = if value.is_a?(Definition)
+              value.to_h
+            elsif (ref_info = refs[camel])
+              wrap_ref(value, ref_info)
+            else
+              value
+            end
           end
         end
       end
@@ -59,6 +66,10 @@ module A2A
         raise "A2A::Schema::Definition should NOT be instantiated directly"
       end
 
+      def self.property_refs
+        raise "A2A::Schema::Definition should NOT be instantiated directly"
+      end
+
       # --- validation ----------------------------------------------------
 
       def valid?
@@ -94,6 +105,29 @@ module A2A
 
       private
 
+        def wrap_ref(value, ref_info)
+          kind, title = ref_info
+
+          case kind
+          when :object
+            value.is_a?(Hash) ? A2A::Schema[title].new(value) : value
+          when :array
+            if value.is_a?(Array)
+              value.map { |el| el.is_a?(Hash) ? A2A::Schema[title].new(el) : el }
+            else
+              value
+            end
+          when :map
+            if value.is_a?(Hash)
+              value.transform_values { |v| v.is_a?(Hash) ? A2A::Schema[title].new(v) : v }
+            else
+              value
+            end
+          else
+            value
+          end
+        end
+
         def deep_compact(obj)
           case obj
           when Hash

data/lib/a2a/schema.rb

@@ -157,6 +157,7 @@ module A2A
           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] }
 
@@ -165,12 +166,14 @@ module A2A
             @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|
@@ -179,6 +182,51 @@ module A2A
           end
         end
 
+        # Inspect each property's raw schema for $ref pointers and build
+        # a map of { camelKey => [:kind, "Definition Title"] } so that
+        # Definition#initialize can auto-wrap nested Hashes.
+        #
+        # Three patterns:
+        #   :object  — direct $ref        (e.g. task → Task)
+        #   :array   — items.$ref          (e.g. artifacts → [Artifact, ...])
+        #   :map     — additionalProperties.$ref (e.g. securitySchemes → {k => SecurityScheme})
+        def build_property_refs(properties)
+          definitions = raw_schema.fetch("definitions", {})
+          refs = {}
+
+          properties.each do |camel_key, prop_schema|
+            if (ref = prop_schema["$ref"])
+              # Direct $ref — singular nested object
+              title = ref_title_for(ref)
+              if title && definitions.dig(title, "properties")
+                refs[camel_key] = [:object, title]
+              end
+            elsif prop_schema["type"] == "array" && (ref = prop_schema.dig("items", "$ref"))
+              # Array with $ref items
+              title = ref_title_for(ref)
+              if title && definitions.dig(title, "properties")
+                refs[camel_key] = [:array, title]
+              end
+            elsif prop_schema["type"] == "object" && (ref = prop_schema.dig("additionalProperties", "$ref"))
+              # Map with $ref additionalProperties
+              title = ref_title_for(ref)
+              if title && definitions.dig(title, "properties")
+                refs[camel_key] = [:map, title]
+              end
+            end
+          end
+
+          refs
+        end
+
+        # Extract the definition title from an internal $ref pointer.
+        # e.g. "#/definitions/Task%20Status" => "Task Status"
+        def ref_title_for(ref)
+          return nil unless ref.start_with?("#/definitions/")
+
+          URI::DEFAULT_PARSER.unescape(ref.sub("#/definitions/", ""))
+        end
+
         def build_snake_to_camel(camel_keys)
           map = {}
           camel_keys.each do |camel|
@@ -426,4 +474,82 @@ test do
       instance.is_a?(A2A::Schema::Definition).should == true
     end
   end
+
+  it "auto-wraps nested $ref Hash into Definition (object pattern)" do
+    response = schema["Send Message Response"].new(
+      task: {
+        "id" => "task-123",
+        "contextId" => "ctx-456",
+        "status" => {
+          "state" => "TASK_STATE_SUBMITTED",
+          "timestamp" => "2025-01-01T00:00:00Z"
+        }
+      }
+    )
+    response.task.should.be.kind_of(A2A::Schema::Definition)
+    response.task.id.should == "task-123"
+    response.task.context_id.should == "ctx-456"
+    # Deeply nested: Task.status is also auto-wrapped
+    response.task.status.should.be.kind_of(A2A::Schema::Definition)
+    response.task.status.state.should == "TASK_STATE_SUBMITTED"
+  end
+
+  it "auto-wraps nested $ref arrays into Definitions (array pattern)" do
+    task = schema["Task"].new(
+      id: "task-1",
+      context_id: "ctx-1",
+      history: [
+        { "role" => "ROLE_USER", "messageId" => "msg-1", "parts" => [{ "text" => "Hi" }] },
+        { "role" => "ROLE_AGENT", "messageId" => "msg-2", "parts" => [{ "text" => "Hello" }] }
+      ]
+    )
+    task.history.should.be.kind_of(Array)
+    task.history.length.should == 2
+    task.history[0].should.be.kind_of(A2A::Schema::Definition)
+    task.history[0].role.should == "ROLE_USER"
+    task.history[1].role.should == "ROLE_AGENT"
+  end
+
+  it "preserves to_h serialization after auto-wrapping" do
+    response = schema["Send Message Response"].new(
+      task: {
+        "id" => "task-123",
+        "contextId" => "ctx-456",
+        "status" => {
+          "state" => "TASK_STATE_SUBMITTED",
+          "timestamp" => "2025-01-01T00:00:00Z"
+        }
+      }
+    )
+    h = response.to_h
+    h["task"].should.be.kind_of(Hash)
+    h["task"]["id"].should == "task-123"
+    h["task"]["status"].should.be.kind_of(Hash)
+    h["task"]["status"]["state"].should == "TASK_STATE_SUBMITTED"
+  end
+
+  it "skips wrapping for opaque types without properties (Struct, Timestamp)" do
+    task = schema["Task"].new(
+      id: "task-1",
+      context_id: "ctx-1",
+      metadata: { "foo" => "bar" }
+    )
+    # metadata refs Struct which has no properties — should stay a plain Hash
+    task.metadata.should.be.kind_of(Hash)
+    task.metadata["foo"].should == "bar"
+  end
+
+  it "does not re-wrap values that are already Definition instances" do
+    status = schema["Task Status"].new(
+      state: "TASK_STATE_SUBMITTED",
+      timestamp: "2025-01-01T00:00:00Z"
+    )
+    task = schema["Task"].new(
+      id: "task-1",
+      context_id: "ctx-1",
+      status: status
+    )
+    # Passing a Definition instance should serialize it (existing behavior)
+    task.to_h["status"]["state"].should == "TASK_STATE_SUBMITTED"
+  end
 end

data/lib/a2a/{bindings → server/bindings}/json_rpc.rb

@@ -5,8 +5,9 @@ require "a2a"
 require "a2a/sse"
 
 module A2A
-  module Bindings
-    # Rack middleware implementing the A2A JSON-RPC 2.0 protocol binding.
+  class Server
+    module Bindings
+      # Rack middleware implementing the A2A JSON-RPC 2.0 protocol binding.
     #
     # Strips the JSON-RPC envelope from the inbound request, setting
     # env keys for the method name, request id, and parsed params.
@@ -47,11 +48,11 @@ module A2A
         env["a2a.json_rpc_method"] = method
         env["a2a.body"]            = params
 
-        @app.call(env)
+        result = @app.call(env)
 
-        # Check if handler signalled a JSON-RPC error
-        if (err = env["a2a.error"])
-          return error_response(id, err[:code], err[:message], err[:data])
+        # Check if the result is an error object
+        if result.is_a?(A2A::Error)
+          return error_response(id, result.code, result.message, result.error_data)
         end
 
         # Check if handler set up a streaming response.
@@ -60,7 +61,6 @@ module A2A
           return [200, A2A::SSE::Stream.headers, stream]
         end
 
-        result = env["a2a.result"]
         success_response(id, result)
       end
 
@@ -79,12 +79,16 @@ module A2A
            [JSON.generate(jsonrpc: "2.0", id: id, error: err)]]
         end
     end
+    end
   end
 end
 
 test do
+  require "a2a/test_helpers"
+
   server = A2A::Server.new(agent_card: { "name" => "Test" })
-  rack   = Rack::MockRequest.new(server)
+  server.register(A2A::TestHelpers.stub_agent)
+  rack = Rack::MockRequest.new(server)
 
   A2A::Proto.operations.each do |op|
     it "#{op.json_rpc_method} returns valid #{op.response_type}" do

data/lib/a2a/{bindings → server/bindings}/rest.rb

@@ -5,8 +5,9 @@ require "a2a"
 require "a2a/sse"
 
 module A2A
-  module Bindings
-    # Rack middleware implementing the A2A HTTP+JSON/REST protocol binding.
+  class Server
+    module Bindings
+      # Rack middleware implementing the A2A HTTP+JSON/REST protocol binding.
     #
     # Extracts the HTTP verb, path, and request body/params into env keys.
     # Calls downstream. On return, wraps env["a2a.result"] into a REST
@@ -40,11 +41,11 @@ module A2A
 
         env["a2a.body"] = params
 
-        @app.call(env)
+        result = @app.call(env)
 
-        # Check if handler signalled a REST error
-        if (err = env["a2a.error"])
-          return error_response(err[:http_status] || 400, err[:message], err[:data])
+        # Check if the result is an error object
+        if result.is_a?(A2A::Error)
+          return error_response(result.http_status, result.message, result.error_data)
         end
 
         # Check if handler set up a streaming response.
@@ -53,7 +54,6 @@ module A2A
           return [200, A2A::SSE::Stream.headers, stream]
         end
 
-        result = env["a2a.result"]
         success_response(result)
       end
 
@@ -72,12 +72,16 @@ module A2A
            [JSON.generate(body)]]
         end
     end
+    end
   end
 end
 
 test do
+  require "a2a/test_helpers"
+
   server = A2A::Server.new(agent_card: { "name" => "Test" })
-  rack   = Rack::MockRequest.new(server)
+  server.register(A2A::TestHelpers.stub_agent)
+  rack = Rack::MockRequest.new(server)
 
   A2A::Proto.operations.each do |op|
     it "#{op.rest_verb.upcase} #{op.rest_path} returns valid #{op.response_type}" do