a2a-rails 0.1.0 → 0.2.0.rc2

data/lib/a2a/rails/protocol/request_handler.rb

@@ -1,13 +1,16 @@
 # frozen_string_literal: true
 
+require "securerandom"
+
 module A2A
   module Rails
     module Protocol
       class RequestHandler
-        def initialize(dispatcher:, lifecycle:, task_mapper: TaskMapper.new)
+        def initialize(dispatcher:, lifecycle:, task_mapper: TaskMapper.new, task_job: TaskExecutionJob)
           @dispatcher = dispatcher
           @lifecycle = lifecycle
           @task_mapper = task_mapper
+          @task_job = task_job
         end
 
         def call(operation:, params:)
@@ -34,22 +37,96 @@ module A2A
           history_length = history_length(params)
           message, context_id = normalize_message(params)
 
+          # Direct Message is an explicit host-side choice, never inferred
+          # from an untrusted client flag or forced by the protocol adapter.
+          if safe_response_mode(message) == :message
+            return send_direct_message(message, context_id)
+          end
+
           task = @lifecycle.create(message: message, context_id: context_id)
-          @lifecycle.start(task.fetch(:id))
-          context = { task_id: task.fetch(:id), context_id: task.fetch(:context_id) }
+          context = execution_context(task)
 
-          task = begin
-            result = @dispatcher.call(message: message, context: context)
-            @lifecycle.complete(task.fetch(:id), result)
+          plan = begin
+            @dispatcher.plan(message: message, context: context)
           rescue RejectedTask => error
-            @lifecycle.reject(task.fetch(:id), error)
+            task = @lifecycle.reject(task.fetch(:id), error)
+            nil
           rescue StandardError => error
-            @lifecycle.fail(task.fetch(:id), error)
+            task = @lifecycle.fail(task.fetch(:id), error)
+            nil
+          end
+
+          if plan&.async?
+            task = enqueue_async_task(task, plan)
+          elsif plan
+            @lifecycle.start(task.fetch(:id))
+            task = begin
+              result = @dispatcher.execute(plan: plan, message: message, context: context)
+              @lifecycle.complete(task.fetch(:id), result)
+            rescue RejectedTask => error
+              @lifecycle.reject(task.fetch(:id), error)
+            rescue StandardError => error
+              @lifecycle.fail(task.fetch(:id), error)
+            end
           end
 
           { "task" => @task_mapper.dump(task, history_length: history_length, include_artifacts: true) }
         end
 
+        def enqueue_async_task(task, plan)
+          arguments = TaskExecutionArguments.new(
+            task_id: task.fetch(:id),
+            principal_id: @lifecycle.principal_id,
+            agent_class_name: plan.agent_class_name,
+            skill_id: plan.skill_id
+          )
+          job = @task_job.perform_later(**arguments.to_h)
+          enqueue_error = job.respond_to?(:enqueue_error) ? job.enqueue_error : nil
+          successfully_enqueued = if job && job.respond_to?(:successfully_enqueued?)
+            job.successfully_enqueued?
+          else
+            !!job
+          end
+          return task if successfully_enqueued && enqueue_error.nil?
+
+          @lifecycle.fail(
+            task.fetch(:id),
+            enqueue_error || ConfigurationError.new("Task execution job was not enqueued")
+          )
+        rescue StandardError => error
+          @lifecycle.fail(task.fetch(:id), error)
+        end
+
+        def execution_context(task)
+          {
+            task_id: task.fetch(:id),
+            context_id: task.fetch(:context_id),
+            principal_id: @lifecycle.principal_id,
+            idempotency_key: task.fetch(:id)
+          }
+        end
+
+        def safe_response_mode(message)
+          @dispatcher.response_mode(message: message)
+        rescue StandardError
+          # The selector is host application code and may raise with secrets.
+          # No Task exists yet to record the failure as a Task status.
+          raise A2A::InvalidAgentResponseError.new
+        end
+
+        def send_direct_message(message, context_id)
+          context_id = SecureRandom.uuid if context_id.nil? || context_id.empty?
+          result = @dispatcher.call(
+            message: message,
+            context: { task_id: nil, context_id: context_id }
+          )
+          { "message" => @task_mapper.dump_direct_message(result, context_id: context_id) }
+        rescue StandardError
+          # No Task exists to record a FAILED status. Never send application
+          # exception details or unsupported result values to A2A callers.
+          raise A2A::InvalidAgentResponseError.new
+        end
+
         def get_task(params)
           history_length = history_length(params)
           task = @lifecycle.find(required_id(params))
@@ -58,6 +135,7 @@ module A2A
 
         def list_tasks(params)
           history_length = history_length(params)
+          validate_list_filters!(params)
           filters = {}
           filters[:context_id] = params["contextId"] if params.key?("contextId")
           filters[:status] = @task_mapper.internal_state(params["status"]) if params.key?("status")
@@ -103,8 +181,21 @@ module A2A
             raise InvalidRequestError, "nonempty parts are required"
           end
 
-          if raw["taskId"].is_a?(String) && !raw["taskId"].empty?
-            @lifecycle.find(raw["taskId"])
+          task_id = raw["taskId"]
+          if !task_id.nil? && !task_id.is_a?(String)
+            raise InvalidRequestError, "taskId must be a String"
+          end
+          context_id = raw["contextId"]
+          if !context_id.nil? && !context_id.is_a?(String)
+            raise InvalidRequestError, "contextId must be a String"
+          end
+          metadata = raw["metadata"]
+          unless metadata.nil? || metadata.is_a?(Hash)
+            raise InvalidRequestError, "message metadata must be an object"
+          end
+
+          if task_id.is_a?(String) && !task_id.empty?
+            @lifecycle.find(task_id)
             raise TaskContinuationNotSupportedError, "Task continuation is not supported in v0.1"
           end
 
@@ -113,25 +204,45 @@ module A2A
             message_id: message_id,
             role: :user,
             parts: normalized_parts,
-            metadata: copy(raw["metadata"] || {})
+            metadata: copy(metadata || {})
           }
 
-          [message, raw["contextId"]]
+          [message, context_id]
         end
 
         def normalize_text_part(part)
           unless part.is_a?(Hash) && part["text"].is_a?(String) && (part.keys & %w[data raw url]).empty?
             raise ContentTypeNotSupportedError, "Only text Parts are supported in v0.1"
           end
+          media_type = part["mediaType"]
+          unless media_type.nil? || media_type.is_a?(String)
+            raise InvalidRequestError, "mediaType must be a String"
+          end
+          metadata = part["metadata"]
+          unless metadata.nil? || metadata.is_a?(Hash)
+            raise InvalidRequestError, "part metadata must be an object"
+          end
 
           normalized = {
             text: part["text"],
             media_type: part["mediaType"].to_s.empty? ? "text/plain" : part["mediaType"]
           }
-          normalized[:metadata] = copy(part["metadata"]) if part.key?("metadata")
+          normalized[:metadata] = copy(metadata) if metadata
           normalized
         end
 
+        def validate_list_filters!(params)
+          if params.key?("contextId") && !params["contextId"].nil? && !params["contextId"].is_a?(String)
+            raise InvalidRequestError, "contextId must be a String"
+          end
+          if params.key?("pageToken") && !params["pageToken"].is_a?(String)
+            raise InvalidRequestError, "pageToken must be a String"
+          end
+          if params.key?("includeArtifacts") && ![true, false].include?(params["includeArtifacts"])
+            raise InvalidRequestError, "includeArtifacts must be boolean"
+          end
+        end
+
         def required_id(params)
           id = params["id"]
           unless id.is_a?(String) && !id.empty?

data/lib/a2a/rails/protocol/task_mapper.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "securerandom"
+
 module A2A
   module Rails
     module Protocol
@@ -33,6 +35,19 @@ module A2A
           result
         end
 
+        def dump_direct_message(value, context_id:)
+          artifact = Task::ArtifactMapper.new.call(value)
+          unless artifact
+            raise ArtifactMappingError, "Direct Message requires a non-nil Handler result"
+          end
+
+          dump_message(
+            message_id: SecureRandom.uuid,
+            role: :agent,
+            parts: artifact.fetch(:parts)
+          ).merge("contextId" => context_id)
+        end
+
         def internal_state(value)
           return nil if value.nil?
 
@@ -86,6 +101,12 @@ module A2A
           result = {}
           result["text"] = part[:text] if part.key?(:text)
           result["data"] = copy(part[:data]) if part.key?(:data)
+          result["raw"] = part[:raw] if part.key?(:raw)
+          result["url"] = part[:url] if part.key?(:url)
+          result["filename"] = part[:filename] if part.key?(:filename)
+          if part.key?(:raw) || part.key?(:url)
+            result["mediaType"] = part[:media_type] if part.key?(:media_type)
+          end
           result["metadata"] = copy(part[:metadata]) if part.key?(:metadata)
           result
         end

data/lib/a2a/rails/request_guard.rb

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+require "stringio"
+require "json"
+
+module A2A
+  module Rails
+    # Enforces a bounded JSON-RPC request body before application callbacks
+    # or the upstream SDK can read/parse arbitrary request data.
+    module RequestGuard
+      DEFAULT_MAX_BYTES = 1_048_576
+      MAX_CONFIGURABLE_BYTES = 16 * 1_048_576
+
+      class PayloadTooLarge < StandardError; end
+      class UnsupportedMediaType < StandardError; end
+      class InvalidBody < StandardError; end
+      class InvalidConfiguration < StandardError; end
+
+      module_function
+
+      def enforce!(env:, max_bytes:)
+        unless max_bytes.is_a?(Integer) && max_bytes.between?(1, MAX_CONFIGURABLE_BYTES)
+          raise InvalidConfiguration, "max_request_bytes must be between 1 and #{MAX_CONFIGURABLE_BYTES}"
+        end
+
+        media_type = env["CONTENT_TYPE"].to_s.split(";", 2).first.to_s.strip.downcase
+        raise UnsupportedMediaType unless media_type == "application/json"
+
+        encoding = env["HTTP_CONTENT_ENCODING"].to_s.strip.downcase
+        raise UnsupportedMediaType unless encoding.empty? || encoding == "identity"
+
+        length = env["CONTENT_LENGTH"]
+        if length && (!length.is_a?(String) || !length.match?(/\A[0-9]+\z/))
+          raise InvalidBody
+        end
+        raise PayloadTooLarge if length && length.to_i > max_bytes
+
+        input = env["rack.input"]
+        raise InvalidBody unless input.respond_to?(:read)
+
+        # Read one byte beyond the permitted maximum to catch requests without
+        # Content-Length, chunked bodies, and falsely understated lengths.
+        body = input.read(max_bytes + 1)
+        raise InvalidBody unless body.is_a?(String)
+        raise PayloadTooLarge if body.bytesize > max_bytes
+        raise InvalidBody if length && length.to_i != body.bytesize
+
+        validate_sdk_sensitive_shapes!(body)
+
+        # The SDK receives the *bounded* copy, never the original input stream.
+        env["rack.input"] = StringIO.new(body)
+        env["CONTENT_LENGTH"] = body.bytesize.to_s
+        true
+      end
+
+      # The upstream agent2agent 2.0.0 ExtractMessage middleware may iterate
+      # SendMessage Parts before our RequestHandler validates them, and can
+      # raise NoMethodError on a null/non-object Part. Reject those shapes at
+      # the bounded HTTP boundary to ensure malformed requests fail safely.
+      def validate_sdk_sensitive_shapes!(body)
+        payload = JSON.parse(body)
+        return unless payload.is_a?(Hash) && payload["method"] == "SendMessage"
+
+        params = payload["params"]
+        message = params["message"] if params.is_a?(Hash)
+        parts = message["parts"] if message.is_a?(Hash)
+
+        unless parts.is_a?(Array) && parts.all? { |part| part.is_a?(Hash) }
+          raise InvalidBody
+        end
+      rescue JSON::ParserError, JSON::NestingError
+        raise InvalidBody
+      end
+    end
+  end
+end

data/lib/a2a/rails/runtime.rb

@@ -3,7 +3,7 @@
 module A2A
   module Rails
     class Runtime
-      def initialize(store: Task::MemoryStore.new)
+      def initialize(store: nil)
         @store = store
       end
 
@@ -14,21 +14,26 @@ module A2A
         AgentCard::Builder.new(
           agent: agent,
           public_base_url: configuration.normalized_public_base_url,
-          request_base_url: request_base_url
+          request_base_url: request_base_url,
+          security: AgentCard::Security.new(configuration: configuration).fields
         ).call
       end
 
-      def call(env:, request_base_url:)
+      def call(env:, request_base_url:, principal_id: nil)
         configuration = A2A::Rails.configuration
         agent = configuration.resolve_agent
         card = AgentCard::Builder.new(
           agent: agent,
           public_base_url: configuration.normalized_public_base_url,
-          request_base_url: request_base_url
+          request_base_url: request_base_url,
+          security: AgentCard::Security.new(configuration: configuration).fields
         ).call
-        lifecycle = Task::Lifecycle.new(store: @store, logger: configuration.logger)
+        # The controller has already verified this Rack-scoped principal.
+        # Never read identity from A2A JSON-RPC parameters or metadata.
+        principal_id ||= env[Authentication::PRINCIPAL_ENV_KEY]
+        lifecycle = Task::Lifecycle.new(store: task_store, logger: configuration.logger, principal_id: principal_id)
         request_handler = Protocol::RequestHandler.new(
-          dispatcher: Dispatcher.new(agent: agent),
+          dispatcher: Dispatcher.new(agent: agent, configuration: configuration),
           lifecycle: lifecycle
         )
         adapter = Protocol::Agent2AgentAdapter.new(
@@ -38,6 +43,10 @@ module A2A
 
         adapter.call(env)
       end
+
+      def task_store
+        @store ||= A2A::Rails.configuration.resolve_task_store
+      end
     end
   end
 end

data/lib/a2a/rails/skill.rb

@@ -3,9 +3,9 @@
 module A2A
   module Rails
     class Skill
-      attr_reader :id, :name, :description, :tags, :handler, :examples, :input_modes, :output_modes
+      attr_reader :id, :name, :description, :tags, :handler, :examples, :input_modes, :output_modes, :execution_mode
 
-      def initialize(id:, description:, tags:, handler:, name: nil, examples: nil, input_modes: nil, output_modes: nil)
+      def initialize(id:, description:, tags:, handler:, name: nil, examples: nil, input_modes: nil, output_modes: nil, execution_mode: nil)
         @id = normalize_id(id)
         @name = normalize_text(name || humanize_id(@id))
         @description = normalize_optional_text(description)
@@ -14,6 +14,7 @@ module A2A
         @examples = normalize_optional_collection(examples)
         @input_modes = normalize_optional_collection(input_modes)
         @output_modes = normalize_optional_collection(output_modes)
+        @execution_mode = execution_mode
         freeze
       end
 
@@ -34,6 +35,10 @@ module A2A
         validate_optional_strings!(:input_modes, @input_modes, allow_empty: false)
         validate_optional_strings!(:output_modes, @output_modes, allow_empty: false)
 
+        unless @execution_mode.nil? || Configuration::TASK_EXECUTION_MODES.include?(@execution_mode)
+          raise ConfigurationError, "Skill #{@id.inspect} execution_mode must be :sync or :async"
+        end
+
         unless @handler.respond_to?(:call)
           raise InvalidHandlerError, "Handler #{handler_label} must respond to .call"
         end