agent2agent 1.0.8 → 1.1.0

data/lib/a2a/middleware/fetch_task.rb

@@ -0,0 +1,228 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "a2a"
+
+module A2A
+  module Middleware
+    # Loads a task from the store and places it on `env["a2a.task"]`.
+    # Sets `env["a2a.task"]` to nil when the ID is absent or the task
+    # does not exist. Does not raise.
+    #
+    # The task ID is read from the request — by default from `request.id`.
+    # Use `id_field:` to read from a different field, and `from:` to
+    # read from a nested object on the request.
+    #
+    # Usage:
+    #
+    #   on "SendMessage" do
+    #     use A2A::Middleware::FetchTask, store: sqlite_store, id_field: :task_id, from: :message
+    #     respond_with -> (env) {
+    #       existing = env["a2a.task"]  # nil if new task
+    #     }
+    #   end
+    #
+    class FetchTask
+      def initialize(app, store:, id_field: :id, from: nil)
+        @app      = app
+        @store    = store
+        @id_field = id_field
+        @from     = from
+      end
+
+      def call(env)
+        request = env["a2a.request"]
+        source  = @from ? request.public_send(@from) : request
+        id      = source.public_send(@id_field)
+
+        env["a2a.task"] = id.to_s.empty? ? nil : @store.get(id)
+
+        @app.call(env)
+      end
+    end
+
+    # Loads a task from the store and places it on `env["a2a.task"]`.
+    # Raises `A2A::TaskNotFoundError` if the task does not exist.
+    #
+    # The task ID is read from the request — by default from `request.id`.
+    # Use `id_field:` to read from a different field, and `from:` to
+    # read from a nested object on the request.
+    #
+    # Usage:
+    #
+    #   on "GetTask" do
+    #     use A2A::Middleware::FetchTaskOrRaise, store: sqlite_store
+    #     respond_with -> (env) {
+    #       task = env["a2a.task"]
+    #     }
+    #   end
+    #
+    #   on "GetTaskPushNotificationConfig" do
+    #     use A2A::Middleware::FetchTaskOrRaise, store: sqlite_store, id_field: :task_id
+    #     respond_with -> (env) { ... }
+    #   end
+    #
+    class FetchTaskOrRaise
+      def initialize(app, store:, id_field: :id, from: nil)
+        @app      = app
+        @store    = store
+        @id_field = id_field
+        @from     = from
+      end
+
+      def call(env)
+        request = env["a2a.request"]
+        source  = @from ? request.public_send(@from) : request
+        id      = source.public_send(@id_field)
+
+        task = @store.get(id)
+        raise A2A::TaskNotFoundError.new(id) unless task
+
+        env["a2a.task"] = task
+
+        @app.call(env)
+      end
+    end
+  end
+end
+
+test do
+  describe "A2A::Middleware::FetchTask" do
+    before do
+      @store = Object.new
+    end
+
+    it "sets env[\"a2a.task\"] when task exists" do
+      task_data = { id: "t-1", state: "TASK_STATE_COMPLETED" }
+      @store.define_singleton_method(:get) { |id| id == "t-1" ? task_data : nil }
+
+      request = Object.new
+      request.define_singleton_method(:id) { "t-1" }
+
+      downstream = -> (env) { env["a2a.task"] }
+      mw = A2A::Middleware::FetchTask.new(downstream, store: @store)
+      env = { "a2a.request" => request }
+
+      mw.call(env).should == task_data
+    end
+
+    it "sets env[\"a2a.task\"] to nil when task does not exist" do
+      @store.define_singleton_method(:get) { |id| nil }
+
+      request = Object.new
+      request.define_singleton_method(:id) { "missing" }
+
+      downstream = -> (env) { env["a2a.task"] }
+      mw = A2A::Middleware::FetchTask.new(downstream, store: @store)
+      env = { "a2a.request" => request }
+
+      mw.call(env).should.be.nil
+    end
+
+    it "sets env[\"a2a.task\"] to nil when id is empty" do
+      request = Object.new
+      request.define_singleton_method(:id) { "" }
+
+      downstream = -> (env) { env["a2a.task"] }
+      mw = A2A::Middleware::FetchTask.new(downstream, store: @store)
+      env = { "a2a.request" => request }
+
+      mw.call(env).should.be.nil
+    end
+
+    it "sets env[\"a2a.task\"] to nil when id is nil" do
+      message = Object.new
+      message.define_singleton_method(:task_id) { nil }
+
+      request = Object.new
+      request.define_singleton_method(:message) { message }
+
+      downstream = -> (env) { env["a2a.task"] }
+      mw = A2A::Middleware::FetchTask.new(downstream, store: @store, id_field: :task_id, from: :message)
+      env = { "a2a.request" => request }
+
+      mw.call(env).should.be.nil
+    end
+
+    it "reads id from a nested object via from:" do
+      task_data = { id: "t-2", state: "TASK_STATE_WORKING" }
+      @store.define_singleton_method(:get) { |id| id == "t-2" ? task_data : nil }
+
+      message = Object.new
+      message.define_singleton_method(:task_id) { "t-2" }
+
+      request = Object.new
+      request.define_singleton_method(:message) { message }
+
+      downstream = -> (env) { env["a2a.task"] }
+      mw = A2A::Middleware::FetchTask.new(downstream, store: @store, id_field: :task_id, from: :message)
+      env = { "a2a.request" => request }
+
+      mw.call(env).should == task_data
+    end
+  end
+
+  describe "A2A::Middleware::FetchTaskOrRaise" do
+    before do
+      @store = Object.new
+    end
+
+    it "sets env[\"a2a.task\"] when task exists" do
+      task_data = { id: "t-1", state: "TASK_STATE_COMPLETED" }
+      @store.define_singleton_method(:get) { |id| id == "t-1" ? task_data : nil }
+
+      request = Object.new
+      request.define_singleton_method(:id) { "t-1" }
+
+      downstream = -> (env) { env["a2a.task"] }
+      mw = A2A::Middleware::FetchTaskOrRaise.new(downstream, store: @store)
+      env = { "a2a.request" => request }
+
+      mw.call(env).should == task_data
+    end
+
+    it "raises TaskNotFoundError when task does not exist" do
+      @store.define_singleton_method(:get) { |id| nil }
+
+      request = Object.new
+      request.define_singleton_method(:id) { "missing" }
+
+      downstream = -> (env) { :should_not_reach }
+      mw = A2A::Middleware::FetchTaskOrRaise.new(downstream, store: @store)
+      env = { "a2a.request" => request }
+
+      lambda { mw.call(env) }.should.raise(A2A::TaskNotFoundError)
+    end
+
+    it "reads task_id when id_field: :task_id is specified" do
+      task_data = { id: "t-2", state: "TASK_STATE_WORKING" }
+      @store.define_singleton_method(:get) { |id| id == "t-2" ? task_data : nil }
+
+      request = Object.new
+      request.define_singleton_method(:task_id) { "t-2" }
+
+      downstream = -> (env) { env["a2a.task"] }
+      mw = A2A::Middleware::FetchTaskOrRaise.new(downstream, store: @store, id_field: :task_id)
+      env = { "a2a.request" => request }
+
+      mw.call(env).should == task_data
+    end
+
+    it "reads id from a nested object via from:" do
+      task_data = { id: "t-3", state: "TASK_STATE_WORKING" }
+      @store.define_singleton_method(:get) { |id| id == "t-3" ? task_data : nil }
+
+      message = Object.new
+      message.define_singleton_method(:task_id) { "t-3" }
+
+      request = Object.new
+      request.define_singleton_method(:message) { message }
+
+      downstream = -> (env) { env["a2a.task"] }
+      mw = A2A::Middleware::FetchTaskOrRaise.new(downstream, store: @store, id_field: :task_id, from: :message)
+      env = { "a2a.request" => request }
+
+      mw.call(env).should == task_data
+    end
+  end
+end

data/lib/a2a/middleware/limit_history_length.rb

@@ -0,0 +1,123 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "a2a"
+
+module A2A
+  module Middleware
+    # Resolves the effective history length limit and sets
+    # `env["a2a.history_length"]` to an integer.
+    #
+    # The server max is required. The effective limit is the *minimum*
+    # of the client-requested value and the server cap. When the client
+    # doesn't specify a history_length, the server cap is used.
+    #
+    # `env["a2a.history_length"]` is always an Integer (0..max).
+    # The handler applies it unconditionally:
+    #
+    #   result["history"] = task[:history]&.last(limit)
+    #
+    # Usage:
+    #
+    #   on "GetTask" do
+    #     use A2A::Middleware::FetchTaskOrRaise, store: sqlite_store
+    #     use A2A::Middleware::LimitHistoryLength, 20
+    #     respond_with -> (env) {
+    #       task  = env["a2a.task"]
+    #       limit = env["a2a.history_length"]
+    #       # ...
+    #       result["history"] = task[:history]&.last(limit)
+    #     }
+    #   end
+    #
+    class LimitHistoryLength
+      def initialize(app, max)
+        @app = app
+        @max = max
+      end
+
+      def call(env)
+        request = env["a2a.request"]
+        limit = @max
+
+        if request.respond_to?(:history_length) && request.history_length
+          limit = [request.history_length.to_i, @max].min
+        end
+
+        env["a2a.history_length"] = limit
+
+        @app.call(env)
+      end
+    end
+  end
+end
+
+test do
+  describe "A2A::Middleware::LimitHistoryLength" do
+    it "defaults to server max when client does not specify" do
+      request = Object.new
+      request.define_singleton_method(:history_length) { nil }
+
+      downstream = -> (env) { env["a2a.history_length"] }
+      mw = A2A::Middleware::LimitHistoryLength.new(downstream, 20)
+      env = { "a2a.request" => request }
+
+      mw.call(env).should == 20
+    end
+
+    it "uses client value when smaller than max" do
+      request = Object.new
+      request.define_singleton_method(:history_length) { 5 }
+
+      downstream = -> (env) { env["a2a.history_length"] }
+      mw = A2A::Middleware::LimitHistoryLength.new(downstream, 20)
+      env = { "a2a.request" => request }
+
+      mw.call(env).should == 5
+    end
+
+    it "caps client value to server max" do
+      request = Object.new
+      request.define_singleton_method(:history_length) { 100 }
+
+      downstream = -> (env) { env["a2a.history_length"] }
+      mw = A2A::Middleware::LimitHistoryLength.new(downstream, 20)
+      env = { "a2a.request" => request }
+
+      mw.call(env).should == 20
+    end
+
+    it "returns 0 when client requests 0" do
+      request = Object.new
+      request.define_singleton_method(:history_length) { 0 }
+
+      downstream = -> (env) { env["a2a.history_length"] }
+      mw = A2A::Middleware::LimitHistoryLength.new(downstream, 20)
+      env = { "a2a.request" => request }
+
+      mw.call(env).should == 0
+    end
+
+    it "handles string values" do
+      request = Object.new
+      request.define_singleton_method(:history_length) { "3" }
+
+      downstream = -> (env) { env["a2a.history_length"] }
+      mw = A2A::Middleware::LimitHistoryLength.new(downstream, 20)
+      env = { "a2a.request" => request }
+
+      mw.call(env).should == 3
+    end
+
+    it "always returns an integer" do
+      request = Object.new
+      request.define_singleton_method(:history_length) { nil }
+
+      downstream = -> (env) { env["a2a.history_length"] }
+      mw = A2A::Middleware::LimitHistoryLength.new(downstream, 10)
+      env = { "a2a.request" => request }
+
+      mw.call(env).should.be.kind_of(Integer)
+    end
+  end
+end

data/lib/a2a/middleware/limit_pagination_size.rb

@@ -0,0 +1,133 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "a2a"
+
+module A2A
+  module Middleware
+    # Clamps `request.page_size` to a valid range and sets
+    # `env["a2a.page_size"]` for downstream handlers.
+    #
+    # Accepts a single integer — the maximum page size (also used as the
+    # default when the client doesn't specify one). Clamps to [1, max].
+    #
+    # Usage:
+    #
+    #   on "ListTasks" do
+    #     use A2A::Middleware::LimitPaginationSize, 50
+    #     respond_with -> (env) {
+    #       page_size = env["a2a.page_size"]
+    #       # ...
+    #     }
+    #   end
+    #
+    class LimitPaginationSize
+      def initialize(app, max = 100)
+        @app = app
+        @max = max
+      end
+
+      def call(env)
+        request = env["a2a.request"]
+
+        page_size = @max
+        if request.respond_to?(:page_size) && request.page_size
+          ps = request.page_size.to_i
+          page_size = [[ps, 1].max, @max].min
+        end
+
+        env["a2a.page_size"] = page_size
+
+        @app.call(env)
+      end
+    end
+  end
+end
+
+test do
+  describe "A2A::Middleware::LimitPaginationSize" do
+    it "uses max as default page size when not specified by client" do
+      request = Object.new
+      request.define_singleton_method(:page_size) { nil }
+
+      downstream = -> (env) { env["a2a.page_size"] }
+      mw = A2A::Middleware::LimitPaginationSize.new(downstream)
+      env = { "a2a.request" => request }
+
+      result = mw.call(env)
+      result.should == 100
+    end
+
+    it "uses a custom max as the default" do
+      request = Object.new
+      request.define_singleton_method(:page_size) { nil }
+
+      downstream = -> (env) { env["a2a.page_size"] }
+      mw = A2A::Middleware::LimitPaginationSize.new(downstream, 25)
+      env = { "a2a.request" => request }
+
+      result = mw.call(env)
+      result.should == 25
+    end
+
+    it "clamps to minimum of 1" do
+      request = Object.new
+      request.define_singleton_method(:page_size) { 0 }
+
+      downstream = -> (env) { env["a2a.page_size"] }
+      mw = A2A::Middleware::LimitPaginationSize.new(downstream)
+      env = { "a2a.request" => request }
+
+      result = mw.call(env)
+      result.should == 1
+    end
+
+    it "clamps to the max" do
+      request = Object.new
+      request.define_singleton_method(:page_size) { 999 }
+
+      downstream = -> (env) { env["a2a.page_size"] }
+      mw = A2A::Middleware::LimitPaginationSize.new(downstream)
+      env = { "a2a.request" => request }
+
+      result = mw.call(env)
+      result.should == 100
+    end
+
+    it "clamps to a custom max" do
+      request = Object.new
+      request.define_singleton_method(:page_size) { 999 }
+
+      downstream = -> (env) { env["a2a.page_size"] }
+      mw = A2A::Middleware::LimitPaginationSize.new(downstream, 50)
+      env = { "a2a.request" => request }
+
+      result = mw.call(env)
+      result.should == 50
+    end
+
+    it "passes through a valid page size" do
+      request = Object.new
+      request.define_singleton_method(:page_size) { 30 }
+
+      downstream = -> (env) { env["a2a.page_size"] }
+      mw = A2A::Middleware::LimitPaginationSize.new(downstream)
+      env = { "a2a.request" => request }
+
+      result = mw.call(env)
+      result.should == 30
+    end
+
+    it "handles string values" do
+      request = Object.new
+      request.define_singleton_method(:page_size) { "20" }
+
+      downstream = -> (env) { env["a2a.page_size"] }
+      mw = A2A::Middleware::LimitPaginationSize.new(downstream)
+      env = { "a2a.request" => request }
+
+      result = mw.call(env)
+      result.should == 20
+    end
+  end
+end