agent2agent 1.1.1 → 2.0.0

data/lib/a2a/server/middleware/sse_stream.rb

@@ -0,0 +1,236 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "a2a"
+require "a2a/server/sse"
+require "async"
+
+module A2A
+  class Server
+    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 agent never calls `open`, the builder is removed from env
+      # so the binding layer doesn't mistake it for a real stream.
+      #
+      # Part of the A2A::Server middleware stack, so the agent gets
+      # env["a2a.stream"] without declaring anything.
+      #
+      # Usage:
+      #
+      #   case env["a2a.operation"]
+      #   in "SendStreamingMessage"
+      #     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::Server::SSE::Stream] the opened stream
+        # @return [nil]
+        #
+        def open(task_id:, context_id:, &block)
+          stream = if @env["a2a.json_rpc_id"]
+            A2A::Server::SSE::JsonRpcStream.new(
+              task_id: task_id, context_id: context_id,
+              json_rpc_id: @env["a2a.json_rpc_id"]
+            )
+          else
+            A2A::Server::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
+end
+
+__END__
+  describe "A2A::Server::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::Server::Middleware::SSEStream.new(downstream)
+      env = {}
+      mw.call(env)
+
+      seen_stream.should.be.kind_of(A2A::Server::Middleware::StreamBuilder)
+    end
+
+    it "auto-clears builder if open was never called" do
+      downstream = ->(env) { :ok }
+
+      mw = A2A::Server::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::Server::Middleware::SSEStream.new(downstream)
+      env = {}
+      mw.call(env)
+
+      env["a2a.stream"].should.be.kind_of(A2A::Server::SSE::RestStream)
+    end
+
+    it "returns the downstream result" do
+      downstream = ->(env) { :result_value }
+
+      mw = A2A::Server::Middleware::SSEStream.new(downstream)
+      result = mw.call({})
+
+      result.should == :result_value
+    end
+  end
+
+  describe "A2A::Server::Middleware::StreamBuilder" do
+    it "creates a RestStream when no JSON-RPC ID" do
+      env = {}
+      builder = A2A::Server::Middleware::StreamBuilder.new(env)
+
+      builder.open(task_id: "t1", context_id: "c1") do |s|
+        s.should.be.kind_of(A2A::Server::SSE::RestStream)
+      end
+
+      env["a2a.stream"].should.be.kind_of(A2A::Server::SSE::RestStream)
+    end
+
+    it "creates a JsonRpcStream when JSON-RPC ID is present" do
+      env = { "a2a.json_rpc_id" => 42 }
+      builder = A2A::Server::Middleware::StreamBuilder.new(env)
+
+      builder.open(task_id: "t1", context_id: "c1") do |s|
+        s.should.be.kind_of(A2A::Server::SSE::JsonRpcStream)
+      end
+
+      env["a2a.stream"].should.be.kind_of(A2A::Server::SSE::JsonRpcStream)
+    end
+
+    it "passes task_id and context_id to the stream" do
+      env = {}
+      builder = A2A::Server::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::Server::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::Server::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::Server::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::Server::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

data/lib/a2a/{sse → server/sse}/event_parser.rb

@@ -3,80 +3,82 @@
 require "json"
 
 module A2A
-  module SSE
-    # Parses raw SSE text chunks into typed StreamResponse schema objects.
-    #
-    # Faraday's `on_data` callback delivers raw text in arbitrary-sized
-    # chunks. This parser buffers them, extracts complete SSE events
-    # (delimited by blank lines), parses the JSON from `data:` lines,
-    # and yields A2A::Schema["Stream Response"] instances.
-    #
-    # For the JSON-RPC binding, the parser unwraps the JSON-RPC 2.0
-    # envelope (`{"jsonrpc":"2.0","id":N,"result":{...}}`) to extract
-    # the StreamResponse payload from `result`.
-    #
-    # Usage:
-    #
-    #   parser = A2A::SSE::EventParser.new(binding: :rest)
-    #
-    #   parser.feed("data: {\"task\":{\"id\":\"t1\"}}\n\n") do |event|
-    #     event.task.id  #=> "t1"
-    #   end
-    #
-    class EventParser
-      def initialize(binding:)
-        @binding = binding
-        @buffer  = String.new
-      end
-
-      # Feed a raw chunk of SSE text. Yields a StreamResponse for each
-      # complete event found in the buffer.
+  class Server
+    module SSE
+      # Parses raw SSE text chunks into typed StreamResponse schema objects.
+      #
+      # Faraday's `on_data` callback delivers raw text in arbitrary-sized
+      # chunks. This parser buffers them, extracts complete SSE events
+      # (delimited by blank lines), parses the JSON from `data:` lines,
+      # and yields A2A::Protocol::JsonSchema["Stream Response"] instances.
+      #
+      # For the JSON-RPC binding, the parser unwraps the JSON-RPC 2.0
+      # envelope (`{"jsonrpc":"2.0","id":N,"result":{...}}`) to extract
+      # the StreamResponse payload from `result`.
       #
-      # @param chunk [String] raw SSE text from the wire
-      # @yieldparam event [A2A::Schema::Definition] a Stream Response instance
+      # Usage:
       #
-      def feed(chunk)
-        @buffer << chunk
-
-        # SSE events are terminated by a blank line (\n\n).
-        # We split on that boundary, keeping any incomplete trailing data
-        # in the buffer for the next call.
-        while (idx = @buffer.index("\n\n"))
-          raw_event = @buffer.slice!(0, idx + 2)
-          payload = parse_event(raw_event)
-          yield A2A::Schema["Stream Response"].new(payload) if payload
+      #   parser = A2A::Server::SSE::EventParser.new(binding: :rest)
+      #
+      #   parser.feed("data: {\"task\":{\"id\":\"t1\"}}\n\n") do |event|
+      #     event.task.id  #=> "t1"
+      #   end
+      #
+      class EventParser
+        def initialize(binding:)
+          @binding = binding
+          @buffer  = String.new
         end
-      end
 
-      private
+        # Feed a raw chunk of SSE text. Yields a StreamResponse for each
+        # complete event found in the buffer.
+        #
+        # @param chunk [String] raw SSE text from the wire
+        # @yieldparam event [A2A::Protocol::JsonSchema::Definition] a Stream Response instance
+        #
+        def feed(chunk)
+          @buffer << chunk
+
+          # SSE events are terminated by a blank line (\n\n).
+          # We split on that boundary, keeping any incomplete trailing data
+          # in the buffer for the next call.
+          while (idx = @buffer.index("\n\n"))
+            raw_event = @buffer.slice!(0, idx + 2)
+            payload = parse_event(raw_event)
+            yield A2A::Protocol::JsonSchema["Stream Response"].new(payload) if payload
+          end
+        end
+
+        private
 
-        # Extract JSON from `data:` lines and optionally unwrap JSON-RPC.
-        def parse_event(raw)
-          data_lines = raw.each_line
-            .select { |line| line.start_with?("data:") }
-            .map { |line| line.sub(/\Adata:\s?/, "").chomp }
+          # Extract JSON from `data:` lines and optionally unwrap JSON-RPC.
+          def parse_event(raw)
+            data_lines = raw.each_line
+              .select { |line| line.start_with?("data:") }
+              .map { |line| line.sub(/\Adata:\s?/, "").chomp }
 
-          return nil if data_lines.empty?
+            return nil if data_lines.empty?
 
-          parsed = JSON.parse(data_lines.join("\n"))
+            parsed = JSON.parse(data_lines.join("\n"))
 
-          case @binding
-          when :json_rpc then parsed["result"]
-          else parsed
+            case @binding
+            when :json_rpc then parsed["result"]
+            else parsed
+            end
+          rescue JSON::ParserError
+            nil
           end
-        rescue JSON::ParserError
-          nil
-        end
+      end
     end
-  end
+end
 end
 
-test do
-  describe "A2A::SSE::EventParser" do
+__END__
+  describe "A2A::Server::SSE::EventParser" do
     # --- REST binding ---
 
     it "parses a single complete SSE event" do
-      parser = A2A::SSE::EventParser.new(binding: :rest)
+      parser = A2A::Server::SSE::EventParser.new(binding: :rest)
       events = []
 
       parser.feed("data: {\"task\":{\"id\":\"t1\",\"contextId\":\"c1\"}}\n\n") do |event|
@@ -84,13 +86,13 @@ test do
       end
 
       events.size.should == 1
-      events[0].should.be.kind_of(A2A::Schema::Definition)
+      events[0].should.be.kind_of(A2A::Protocol::JsonSchema::Definition)
       events[0].task.should.not.be.nil
       events[0].task.id.should == "t1"
     end
 
     it "parses multiple events in a single chunk" do
-      parser = A2A::SSE::EventParser.new(binding: :rest)
+      parser = A2A::Server::SSE::EventParser.new(binding: :rest)
       events = []
 
       chunk = [
@@ -106,7 +108,7 @@ test do
     end
 
     it "handles partial chunks via buffering" do
-      parser = A2A::SSE::EventParser.new(binding: :rest)
+      parser = A2A::Server::SSE::EventParser.new(binding: :rest)
       events = []
 
       # First chunk: incomplete event
@@ -120,7 +122,7 @@ test do
     end
 
     it "handles events with event: and id: fields" do
-      parser = A2A::SSE::EventParser.new(binding: :rest)
+      parser = A2A::Server::SSE::EventParser.new(binding: :rest)
       events = []
 
       chunk = "event: message\nid: 42\ndata: {\"task\":{\"id\":\"t1\",\"contextId\":\"c1\"}}\n\n"
@@ -131,7 +133,7 @@ test do
     end
 
     it "skips events with no data lines" do
-      parser = A2A::SSE::EventParser.new(binding: :rest)
+      parser = A2A::Server::SSE::EventParser.new(binding: :rest)
       events = []
 
       parser.feed(": comment\n\n") { |e| events << e }
@@ -139,7 +141,7 @@ test do
     end
 
     it "skips events with invalid JSON" do
-      parser = A2A::SSE::EventParser.new(binding: :rest)
+      parser = A2A::Server::SSE::EventParser.new(binding: :rest)
       events = []
 
       parser.feed("data: not-json\n\n") { |e| events << e }
@@ -149,7 +151,7 @@ test do
     # --- JSON-RPC binding ---
 
     it "unwraps JSON-RPC 2.0 envelope" do
-      parser = A2A::SSE::EventParser.new(binding: :json_rpc)
+      parser = A2A::Server::SSE::EventParser.new(binding: :json_rpc)
       events = []
 
       envelope = {
@@ -168,7 +170,7 @@ test do
     end
 
     it "handles multiple JSON-RPC events" do
-      parser = A2A::SSE::EventParser.new(binding: :json_rpc)
+      parser = A2A::Server::SSE::EventParser.new(binding: :json_rpc)
       events = []
 
       [
@@ -186,7 +188,7 @@ test do
     # --- Multi-line data ---
 
     it "handles multi-line data fields" do
-      parser = A2A::SSE::EventParser.new(binding: :rest)
+      parser = A2A::Server::SSE::EventParser.new(binding: :rest)
       events = []
 
       # JSON split across multiple data: lines
@@ -199,4 +201,3 @@ test do
       events[0].task.id.should == "t1"
     end
   end
-end

data/lib/a2a/server/sse/json_rpc_stream.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+require_relative "stream"
+
+module A2A
+  class Server
+    module SSE
+      # SSE stream that wraps each event in a JSON-RPC 2.0 response envelope.
+      #
+      # Per the A2A spec, JSON-RPC streaming returns SSE where each `data:` line
+      # is a full JSON-RPC response: {"jsonrpc":"2.0","id":N,"result":{...}}
+      #
+      # Usage:
+      #
+      #   stream = A2A::Server::SSE::JsonRpcStream.new(
+      #     task_id: "t1", context_id: "c1", json_rpc_id: 1
+      #   )
+      #
+      #   Async do
+      #     stream.task(status: { state: "TASK_STATE_WORKING", timestamp: "..." })
+      #     stream.finish
+      #   end
+      #
+      #   [200, A2A::Server::SSE::Stream.headers, stream]
+      #
+      class JsonRpcStream < Stream
+        def initialize(json_rpc_id:, **options)
+          @json_rpc_id = json_rpc_id
+          super(**options)
+        end
+
+        # Emit an SSE event wrapped in a JSON-RPC envelope.
+        #
+        # @param data [Hash] the StreamResponse payload (becomes "result")
+        #
+        def event(data, **opts)
+          envelope = {
+            "jsonrpc" => "2.0",
+            "id"      => @json_rpc_id,
+            "result"  => data.respond_to?(:to_h) ? data.to_h : data,
+          }
+          super(envelope, **opts)
+        end
+      end
+    end
+end
+end
+
+__END__
+  describe "A2A::Server::SSE::JsonRpcStream" do
+    it "wraps events in JSON-RPC 2.0 envelopes" do
+      stream = A2A::Server::SSE::JsonRpcStream.new(
+        task_id: "t1", context_id: "c1", json_rpc_id: 42
+      )
+
+      stream.task(status: { state: "TASK_STATE_WORKING", timestamp: "2025-01-01T00:00:00Z" })
+      stream.finish
+
+      chunk = stream.read
+      chunk.should.include('"jsonrpc"')
+      chunk.should.include('"2.0"')
+      chunk.should.include('"id":42')  # numeric id preserved
+      chunk.should.include('"result"')
+    end
+
+    it "is a subclass of SSE::Stream" do
+      stream = A2A::Server::SSE::JsonRpcStream.new(
+        task_id: "t1", context_id: "c1", json_rpc_id: 1
+      )
+      stream.is_a?(A2A::Server::SSE::Stream).should == true
+      stream.is_a?(::Protocol::HTTP::Body::Readable).should == true
+    end
+
+    it "typed methods chain through JSON-RPC envelope" do
+      stream = A2A::Server::SSE::JsonRpcStream.new(
+        task_id: "t1", context_id: "c1", json_rpc_id: 7
+      )
+
+      stream.status_update(status: { state: "TASK_STATE_COMPLETED", timestamp: "2025-01-01T00:00:00Z" })
+      stream.finish
+
+      chunk = stream.read
+      parsed = JSON.parse(chunk.sub(/\Adata: /, "").strip)
+      parsed["jsonrpc"].should == "2.0"
+      parsed["id"].should == 7
+      parsed["result"]["statusUpdate"]["taskId"].should == "t1"
+      parsed["result"]["statusUpdate"]["contextId"].should == "c1"
+    end
+  end

data/lib/a2a/server/sse/rest_stream.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+require_relative "stream"
+
+module A2A
+  class Server
+    module SSE
+      # SSE stream for the HTTP+JSON/REST binding.
+      #
+      # Per the A2A spec, REST streaming returns SSE where each `data:` line
+      # is a bare StreamResponse JSON object (no envelope wrapping).
+      #
+      # This is essentially just an alias for Stream — the base class already
+      # emits bare JSON. We keep this as a named class for symmetry with
+      # JsonRpcStream and to make binding code self-documenting.
+      #
+      # Usage:
+      #
+      #   stream = A2A::Server::SSE::RestStream.new(task_id: "t1", context_id: "c1")
+      #
+      #   Async do
+      #     stream.task(status: { state: "TASK_STATE_WORKING", timestamp: "..." })
+      #     stream.finish
+      #   end
+      #
+      #   [200, A2A::Server::SSE::Stream.headers, stream]
+      #
+      class RestStream < Stream
+      end
+    end
+end
+end
+
+__END__
+  describe "A2A::Server::SSE::RestStream" do
+    it "emits bare JSON events (no envelope)" do
+      stream = A2A::Server::SSE::RestStream.new(task_id: "t1", context_id: "c1")
+
+      stream.status_update(status: { state: "TASK_STATE_COMPLETED", timestamp: "2025-01-01T00:00:00Z" })
+      stream.finish
+
+      chunk = stream.read
+      chunk.should.include('"statusUpdate"')
+      chunk.should.not.include('"jsonrpc"')
+    end
+
+    it "is a subclass of SSE::Stream" do
+      stream = A2A::Server::SSE::RestStream.new(task_id: "t1", context_id: "c1")
+      stream.is_a?(A2A::Server::SSE::Stream).should == true
+    end
+
+    it "injects task_id and context_id into typed events" do
+      stream = A2A::Server::SSE::RestStream.new(task_id: "t1", context_id: "c1")
+
+      stream.task(status: { state: "TASK_STATE_WORKING" })
+      stream.finish
+
+      chunk = stream.read
+      parsed = JSON.parse(chunk.sub(/\Adata: /, "").strip)
+      parsed["task"]["id"].should == "t1"
+      parsed["task"]["contextId"].should == "c1"
+    end
+  end