ai-agents 0.5.0 → 0.6.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: d07ac97ca06177ee4504601099af6bc6ba3e4d8557b28d34cd3222240b27575d
-  data.tar.gz: cd0c7121918c8a28c3760325c9acd319e86d349f34fb715d6b953912a99f98a8
+  metadata.gz: 7f8fa7ec73784bc0fb1e9e6bd4852c6d0295160c4ca28d435c895cba28f5cd58
+  data.tar.gz: 3e489f11ae2a5c93b4232ec78e3c7385c83f303c28039ed35fa082669fedaf4b
 SHA512:
-  metadata.gz: 63db003d8bf43b2ba8d52d48dbb24d2f8bf86ee4b4e8ab704316cf89d0b80156110a8438a7ca86a6fc37182ff9f5a723ace7c474e53a078de050cb22eff4b268
-  data.tar.gz: 05606d58c663ca0b9b062ed537cba33a4b987e5fd720a1188c003150f880a7c953cc28096ca4024854dcdec0a9896a7da35653ebb4f95c2e04a0f5b97e3e1c2e
+  metadata.gz: ba5050ba743466993826d888673d90696d27f422df3d3972027fa141de05d6436c45cd44dba0836dc6350ad83c0ff1628ba43bbce6d93e5e49e8efc450554e91
+  data.tar.gz: efed7a181a6f52c4c8feb8b84e1bbb2309b5216a85569dae69dd4038cde0cdad7a283e066181414e50657aee063e79935573799fdc24abc6253b02bcb8e5d09f

data/.rubocop.yml

@@ -10,20 +10,18 @@ Style/StringLiterals:
 Style/StringLiteralsInInterpolation:
   EnforcedStyle: double_quotes
 
+Metrics/MethodLength:
+  Max: 20
+Metrics/ClassLength:
+  Enabled: false
+
+RSpec/MultipleDescribes:
+  Enabled: false
 RSpec/MultipleExpectations:
   Max: 10
-
 RSpec/ExampleLength:
   Max: 20
-
 RSpec/MultipleMemoizedHelpers:
   Max: 15
-
 RSpec/SpecFilePathFormat:
   Enabled: false
-
-Metrics/MethodLength:
-  Max: 20
-
-RSpec/MultipleDescribes:
-  Enabled: false

data/CHANGELOG.md

@@ -5,6 +5,25 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.6.0] - 2025-10-16
+
+### Added
+- **Custom HTTP Headers Support**: Agents can now specify custom HTTP headers for LLM requests
+  - Added `headers` parameter to `Agent#initialize` for setting agent-level default headers
+  - Runtime headers can be passed via `headers` parameter in `AgentRunner#run` method
+  - Runtime headers take precedence over agent-level headers when keys overlap
+  - Headers are automatically normalized (symbolized keys) and validated
+  - Full support for headers across agent handoffs with proper merging logic
+  - New `Agents::Helpers::Headers` module for header normalization and merging
+  - Comprehensive test coverage for header functionality
+
+### Changed
+- **Code Organization**: Refactored internal helpers into dedicated module structure
+  - Moved `MessageExtractor` to `Agents::Helpers::MessageExtractor` module
+  - Converted `MessageExtractor` from class-based to module-function pattern
+  - Created `lib/agents/helpers/` directory for helper modules
+  - All helper modules now use consistent flat naming convention (`Agents::Helpers::ModuleName`)
+
 ## [0.5.0] - 2025-08-20
 
 ### Added

data/docs/guides/request-headers.md

@@ -0,0 +1,91 @@
+---
+layout: default
+title: Custom Request Headers
+parent: Guides
+nav_order: 6
+---
+
+# Custom Request Headers
+
+Custom HTTP headers allow you to pass additional metadata with your LLM API requests. This is useful for authentication, request tracking, A/B testing, and provider-specific features.
+
+## Basic Usage
+
+### Agent-Level Headers
+
+Set default headers when creating an agent that will be applied to all requests:
+
+```ruby
+agent = Agents::Agent.new(
+  name: "Assistant",
+  instructions: "You are a helpful assistant",
+  headers: {
+    "X-Custom-ID" => "agent-123",
+    "X-Environment" => "production"
+  }
+)
+
+runner = Agents::Runner.with_agents(agent)
+result = runner.run("Hello!")
+# All requests will include the custom headers
+```
+
+### Runtime Headers
+
+Override or add headers for specific requests:
+
+```ruby
+agent = Agents::Agent.new(
+  name: "Assistant",
+  instructions: "You are a helpful assistant"
+)
+
+runner = Agents::Runner.with_agents(agent)
+
+# Pass headers at runtime
+result = runner.run(
+  "What's the weather?",
+  headers: {
+    "X-Request-ID" => "req-456",
+    "X-User-ID" => "user-789"
+  }
+)
+```
+
+### Header precedence
+
+When both agent-level and runtime headers are provided, **runtime headers take precedence**:
+
+```ruby
+agent = Agents::Agent.new(
+  name: "Assistant",
+  instructions: "You are a helpful assistant",
+  headers: {
+    "X-Environment" => "staging",
+    "X-Agent-ID" => "agent-001"
+  }
+)
+
+runner = Agents::Runner.with_agents(agent)
+
+result = runner.run(
+  "Hello!",
+  headers: {
+    "X-Environment" => "production",  # Overrides agent's staging value
+    "X-Request-ID" => "req-123"       # Additional header
+  }
+)
+
+# Final headers sent to LLM API:
+# {
+#   "X-Environment" => "production",  # Runtime value wins
+#   "X-Agent-ID" => "agent-001",      # From agent
+#   "X-Request-ID" => "req-123"       # From runtime
+# }
+```
+
+## See Also
+
+- [Multi-Agent Systems](multi-agent-systems.html) - Using headers across agent handoffs
+- [Rails Integration](rails-integration.html) - Request tracking in Rails applications
+- [State Persistence](state-persistence.html) - Combining headers with conversation state

data/docs/guides.md

@@ -17,3 +17,4 @@ Practical guides for building real-world applications with the AI Agents library
 - **[Rails Integration](guides/rails-integration.html)** - Integrating agents with Ruby on Rails applications and ActiveRecord persistence
 - **[State Persistence](guides/state-persistence.html)** - Managing conversation state and context across sessions and processes
 - **[Structured Output](guides/structured-output.html)** - Enforcing JSON schema validation for reliable agent responses
+- **[Custom Request Headers](guides/request-headers.html)** - Adding custom HTTP headers for authentication, tracking, and provider-specific features

data/examples/isp-support/agents_factory.rb

@@ -6,6 +6,7 @@ require_relative "tools/create_lead_tool"
 require_relative "tools/create_checkout_tool"
 require_relative "tools/search_docs_tool"
 require_relative "tools/escalate_to_human_tool"
+require "ruby_llm/schema"
 
 module ISPSupport
   # Factory for creating all ISP support agents with proper handoff relationships.
@@ -56,7 +57,8 @@ module ISPSupport
         instructions: sales_instructions_with_state,
         model: "gpt-4.1-mini",
         tools: [ISPSupport::CreateLeadTool.new, ISPSupport::CreateCheckoutTool.new],
-        temperature: 0.8  # Higher temperature for more persuasive, varied sales language
+        temperature: 0.8, # Higher temperature for more persuasive, varied sales language
+        response_schema: sales_response_schema
       )
     end
 
@@ -70,7 +72,8 @@ module ISPSupport
           ISPSupport::SearchDocsTool.new,
           ISPSupport::EscalateToHumanTool.new
         ],
-        temperature: 0.5  # Balanced temperature for helpful but consistent technical support
+        temperature: 0.5, # Balanced temperature for helpful but consistent technical support
+        response_schema: triage_response_schema
       )
     end
 
@@ -95,22 +98,33 @@ module ISPSupport
     end
 
     def triage_response_schema
-      {
-        type: "object",
-        properties: {
-          response: {
-            type: "string",
-            description: "Your response to the customer"
-          },
-          intent: {
-            type: "string",
-            enum: %w[sales support unclear],
-            description: "The detected intent category"
-          }
-        },
-        required: %w[response intent],
-        additionalProperties: false
-      }
+      RubyLLM::Schema.create do
+        string :response, description: "Your response to the customer"
+        string :intent, enum: %w[sales support unclear], description: "The detected intent category"
+        array :sentiment, description: "Customer sentiment indicators" do
+          string enum: %w[positive neutral negative frustrated urgent confused satisfied]
+        end
+      end
+    end
+
+    def support_response_schema
+      RubyLLM::Schema.create do
+        string :response, description: "Your response to the customer"
+        string :intent, enum: %w[support], description: "The intent category (always support)"
+        array :sentiment, description: "Customer sentiment indicators" do
+          string enum: %w[positive neutral negative frustrated urgent confused satisfied]
+        end
+      end
+    end
+
+    def sales_response_schema
+      RubyLLM::Schema.create do
+        string :response, description: "Your response to the customer"
+        string :intent, enum: %w[sales], description: "The intent category (always sales)"
+        array :sentiment, description: "Customer sentiment indicators" do
+          string enum: %w[positive neutral negative frustrated urgent confused satisfied]
+        end
+      end
     end
 
     def sales_instructions

data/examples/isp-support/interactive.rb

@@ -2,6 +2,7 @@
 # frozen_string_literal: true
 
 require "json"
+require "readline"
 require_relative "../../lib/agents"
 require_relative "agents_factory"
 
@@ -29,41 +30,59 @@ class ISPSupportDemo
     @context = {}
     @current_status = ""
 
-    puts "🏢 Welcome to ISP Customer Support!"
-    puts "Type '/help' for commands or 'exit' to quit."
+    puts green("🏢 Welcome to ISP Customer Support!")
+    puts dim_text("Type '/help' for commands or 'exit' to quit.")
     puts
   end
 
   def start
     loop do
-      print "💬 You: "
-      user_input = gets.chomp.strip
+      user_input = Readline.readline(cyan("\u{1F4AC} You: "), true)
+      next unless user_input # Handle Ctrl+D
 
+      user_input = user_input.strip
       command_result = handle_command(user_input)
       break if command_result == :exit
       next if command_result == :handled || user_input.empty?
 
       # Clear any previous status and show agent is working
       clear_status_line
-      print "🤖 Processing..."
+      print yellow("🤖 Processing...")
 
-      # Use the runner - it automatically determines the right agent from context
-      result = @runner.run(user_input, context: @context)
+      begin
+        # Use the runner - it automatically determines the right agent from context
+        result = @runner.run(user_input, context: @context)
 
-      # Update our context with the returned context from Runner
-      @context = result.context if result.respond_to?(:context) && result.context
+        # Update our context with the returned context from Runner
+        @context = result.context if result.respond_to?(:context) && result.context
 
-      # Clear status and show response
-      clear_status_line
+        # Clear status and show response with callback history
+        clear_status_line
 
-      # Handle structured output from triage agent
-      output = result.output || "[No output]"
-      if @context[:current_agent] == "Triage Agent" && output.is_a?(Hash)
-        # Display the response from structured response
-        puts "🤖 #{output["response"]}"
-        puts "\e[2m   [Intent]: #{output["intent"]}\e[0m" if output["intent"]
-      else
-        puts "🤖 #{output}"
+        # Display callback messages if any
+        if @callback_messages.any?
+          puts dim_text(@callback_messages.join("\n"))
+          @callback_messages.clear
+        end
+
+        # Handle structured output from agents
+        output = result.output || "[No output]"
+
+        if output.is_a?(Hash) && output.key?("response")
+          # Display the response from structured response
+          puts "🤖 #{output["response"]}"
+          puts dim_text("   [Intent]: #{output["intent"]}") if output["intent"]
+          puts dim_text("   [Sentiment]: #{output["sentiment"].join(", ")}") if output["sentiment"]&.any?
+        else
+          puts "🤖 #{output}"
+        end
+
+        puts # Add blank line after agent response
+      rescue StandardError => e
+        clear_status_line
+        puts red("❌ Error: #{e.message}")
+        puts dim_text("Please try again or type '/help' for assistance.")
+        puts # Add blank line after error message
       end
     end
   end
@@ -71,26 +90,36 @@ class ISPSupportDemo
   private
 
   def setup_callbacks
+    @callback_messages = []
+
     @runner.on_agent_thinking do |agent_name, _input|
-      update_status("🧠 #{agent_name} is thinking...")
+      message = "🧠 #{agent_name} is thinking..."
+      update_status(message)
+      @callback_messages << message
     end
 
     @runner.on_tool_start do |tool_name, _args|
-      update_status("🔧 Using #{tool_name}...")
+      message = "🔧 Using #{tool_name}..."
+      update_status(message)
+      @callback_messages << message
     end
 
     @runner.on_tool_complete do |tool_name, _result|
-      update_status("✅ #{tool_name} completed")
+      message = "✅ #{tool_name} completed"
+      update_status(message)
+      @callback_messages << message
     end
 
     @runner.on_agent_handoff do |from_agent, to_agent, _reason|
-      update_status("🔄 Handoff: #{from_agent} → #{to_agent}")
+      message = "🔄 Handoff: #{from_agent} → #{to_agent}"
+      update_status(message)
+      @callback_messages << message
     end
   end
 
   def update_status(message)
     clear_status_line
-    print message
+    print dim_text(message)
     $stdout.flush
   end
 
@@ -197,6 +226,27 @@ class ISPSupportDemo
     else "Unknown agent"
     end
   end
+
+  # ANSI color helper methods
+  def dim_text(text)
+    "\e[90m#{text}\e[0m"
+  end
+
+  def green(text)
+    "\e[32m#{text}\e[0m"
+  end
+
+  def yellow(text)
+    "\e[33m#{text}\e[0m"
+  end
+
+  def red(text)
+    "\e[31m#{text}\e[0m"
+  end
+
+  def cyan(text)
+    "\e[36m#{text}\e[0m"
+  end
 end
 
 # Run the demo

data/lib/agents/agent.rb

@@ -4,7 +4,7 @@
 # Agents are immutable, thread-safe objects that can be cloned with modifications.
 # They encapsulate the configuration needed to interact with an LLM including
 # instructions, tools, and potential handoff targets.
-#
+require_relative "helpers/headers"
 # @example Creating a basic agent
 #   agent = Agents::Agent.new(
 #     name: "Assistant",
@@ -50,7 +50,7 @@
 #   )
 module Agents
   class Agent
-    attr_reader :name, :instructions, :model, :tools, :handoff_agents, :temperature, :response_schema
+    attr_reader :name, :instructions, :model, :tools, :handoff_agents, :temperature, :response_schema, :headers
 
     # Initialize a new Agent instance
     #
@@ -61,8 +61,9 @@ module Agents
     # @param handoff_agents [Array<Agents::Agent>] Array of agents this agent can hand off to
     # @param temperature [Float] Controls randomness in responses (0.0 = deterministic, 1.0 = very random, default: 0.7)
     # @param response_schema [Hash, nil] JSON schema for structured output responses
+    # @param headers [Hash, nil] Default HTTP headers applied to LLM requests
     def initialize(name:, instructions: nil, model: "gpt-4.1-mini", tools: [], handoff_agents: [], temperature: 0.7,
-                   response_schema: nil)
+                   response_schema: nil, headers: nil)
       @name = name
       @instructions = instructions
       @model = model
@@ -70,6 +71,7 @@ module Agents
       @handoff_agents = []
       @temperature = temperature
       @response_schema = response_schema
+      @headers = Helpers::Headers.normalize(headers, freeze_result: true)
 
       # Mutex for thread-safe handoff registration
       # While agents are typically configured at startup, we want to ensure
@@ -164,7 +166,8 @@ module Agents
         tools: changes.fetch(:tools, @tools.dup),
         handoff_agents: changes.fetch(:handoff_agents, @handoff_agents),
         temperature: changes.fetch(:temperature, @temperature),
-        response_schema: changes.fetch(:response_schema, @response_schema)
+        response_schema: changes.fetch(:response_schema, @response_schema),
+        headers: changes.fetch(:headers, @headers)
       )
     end
 

data/lib/agents/agent_runner.rb

@@ -58,12 +58,12 @@ module Agents
     # @param input [String] User's message
     # @param context [Hash] Conversation context (will be restored if continuing conversation)
     # @param max_turns [Integer] Maximum turns before stopping (default: 10)
+    # @param headers [Hash, nil] Custom HTTP headers to pass through to the underlying LLM provider
     # @return [RunResult] Execution result with output, messages, and updated context
-    def run(input, context: {}, max_turns: Runner::DEFAULT_MAX_TURNS)
+    def run(input, context: {}, max_turns: Runner::DEFAULT_MAX_TURNS, headers: nil)
       # Determine which agent should handle this conversation
       # Uses conversation history to maintain continuity across handoffs
       current_agent = determine_conversation_agent(context)
-
       # Execute using stateless Runner - each execution is independent and thread-safe
       # Pass callbacks to enable real-time event notifications
       Runner.new.run(
@@ -72,6 +72,7 @@ module Agents
         context: context,
         registry: @registry,
         max_turns: max_turns,
+        headers: headers,
         callbacks: @callbacks
       )
     end

data/lib/agents/helpers/headers.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Agents::Helpers::Headers
+  module_function
+
+  def normalize(headers, freeze_result: false)
+    return freeze_result ? {}.freeze : {} if headers.nil? || (headers.respond_to?(:empty?) && headers.empty?)
+
+    hash = headers.respond_to?(:to_h) ? headers.to_h : headers
+    raise ArgumentError, "headers must be a Hash or respond to #to_h" unless hash.is_a?(Hash)
+
+    result = symbolize_keys(hash)
+    freeze_result ? result.freeze : result
+  end
+
+  def merge(agent_headers, runtime_headers)
+    return runtime_headers if agent_headers.empty?
+    return agent_headers if runtime_headers.empty?
+
+    agent_headers.merge(runtime_headers) { |_key, _agent_value, runtime_value| runtime_value }
+  end
+
+  def symbolize_keys(hash)
+    hash.each_with_object({}) do |(key, value), memo|
+      memo[key.is_a?(Symbol) ? key : key.to_sym] = value
+    end
+  end
+  private_class_method :symbolize_keys
+end

data/lib/agents/helpers/message_extractor.rb

@@ -0,0 +1,88 @@
+# frozen_string_literal: true
+
+# Service object responsible for extracting and formatting conversation messages
+# from RubyLLM chat objects into a format suitable for persistence and context restoration.
+#
+# Handles different message types:
+# - User messages: Basic content preservation
+# - Assistant messages: Includes agent attribution and tool calls
+# - Tool result messages: Links back to original tool calls
+#
+# @example Extract messages from a chat
+#   messages = Agents::Helpers::MessageExtractor.extract_messages(chat, current_agent)
+#   #=> [
+#     { role: :user, content: "Hello" },
+#     { role: :assistant, content: "Hi!", agent_name: "Support", tool_calls: [...] },
+#     { role: :tool, content: "Result", tool_call_id: "call_123" }
+#   ]
+module Agents::Helpers::MessageExtractor
+  module_function
+
+  # Check if content is considered empty (handles both String and Hash content)
+  #
+  # @param content [String, Hash, nil] The content to check
+  # @return [Boolean] true if content is empty, false otherwise
+  def content_empty?(content)
+    case content
+    when String
+      content.strip.empty?
+    when Hash
+      content.empty?
+    else
+      content.nil?
+    end
+  end
+
+  # Extract messages from a chat object for conversation history persistence
+  #
+  # @param chat [Object] Chat object that responds to :messages
+  # @param current_agent [Agent] The agent currently handling the conversation
+  # @return [Array<Hash>] Array of message hashes suitable for persistence
+  def extract_messages(chat, current_agent)
+    return [] unless chat.respond_to?(:messages)
+
+    chat.messages.filter_map do |msg|
+      case msg.role
+      when :user, :assistant
+        extract_user_or_assistant_message(msg, current_agent)
+      when :tool
+        extract_tool_message(msg)
+      end
+    end
+  end
+
+  def extract_user_or_assistant_message(msg, current_agent)
+    return nil unless msg.content && !content_empty?(msg.content)
+
+    message = {
+      role: msg.role,
+      content: msg.content
+    }
+
+    if msg.role == :assistant
+      # Add agent attribution for conversation continuity
+      message[:agent_name] = current_agent.name if current_agent
+
+      # Add tool calls if present
+      if msg.tool_call? && msg.tool_calls
+        # RubyLLM stores tool_calls as Hash with call_id => ToolCall object
+        # Reference: RubyLLM::StreamAccumulator#tool_calls_from_stream
+        message[:tool_calls] = msg.tool_calls.values.map(&:to_h)
+      end
+    end
+
+    message
+  end
+  private_class_method :extract_user_or_assistant_message
+
+  def extract_tool_message(msg)
+    return nil unless msg.tool_result?
+
+    {
+      role: msg.role,
+      content: msg.content,
+      tool_call_id: msg.tool_call_id
+    }
+  end
+  private_class_method :extract_tool_message
+end

data/lib/agents/helpers.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module Agents
+  module Helpers
+  end
+end
+
+require_relative "helpers/headers"
+require_relative "helpers/message_extractor"

data/lib/agents/runner.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require_relative "message_extractor"
-
 module Agents
   # The execution engine that orchestrates conversations between users and agents.
   # Runner manages the conversation flow, handles tool execution through RubyLLM,
@@ -55,6 +53,7 @@ module Agents
     DEFAULT_MAX_TURNS = 10
 
     class MaxTurnsExceeded < StandardError; end
+    class AgentNotFoundError < StandardError; end
 
     # Create a thread-safe agent runner for multi-agent conversations.
     # The first agent becomes the default entry point for new conversations.
@@ -79,9 +78,10 @@ module Agents
     # @param context [Hash] Shared context data accessible to all tools
     # @param registry [Hash] Registry of agents for handoff resolution
     # @param max_turns [Integer] Maximum conversation turns before stopping
+    # @param headers [Hash, nil] Custom HTTP headers passed to the underlying LLM provider
     # @param callbacks [Hash] Optional callbacks for real-time event notifications
     # @return [RunResult] The result containing output, messages, and usage
-    def run(starting_agent, input, context: {}, registry: {}, max_turns: DEFAULT_MAX_TURNS, callbacks: {})
+    def run(starting_agent, input, context: {}, registry: {}, max_turns: DEFAULT_MAX_TURNS, headers: nil, callbacks: {})
       # The starting_agent is already determined by AgentRunner based on conversation history
       current_agent = starting_agent
 
@@ -90,15 +90,22 @@ module Agents
       context_wrapper = RunContext.new(context_copy, callbacks: callbacks)
       current_turn = 0
 
+      runtime_headers = Helpers::Headers.normalize(headers)
+      agent_headers = Helpers::Headers.normalize(current_agent.headers)
+
       # Create chat and restore conversation history
-      chat = create_chat(current_agent, context_wrapper)
+      chat = RubyLLM::Chat.new(model: current_agent.model)
+      current_headers = Helpers::Headers.merge(agent_headers, runtime_headers)
+      apply_headers(chat, current_headers)
+      configure_chat_for_agent(chat, current_agent, context_wrapper, replace: false)
       restore_conversation_history(chat, context_wrapper)
 
+
       loop do
         current_turn += 1
         raise MaxTurnsExceeded, "Exceeded maximum turns: #{max_turns}" if current_turn > max_turns
 
-        # Get response from LLM (Extended Chat handles tool execution with handoff detection)
+        # Get response from LLM (RubyLLM handles tool execution with halting based handoff detection)
         result = if current_turn == 1
                    # Emit agent thinking event for initial message
                    context_wrapper.callback_manager.emit_agent_thinking(current_agent.name, input)
@@ -118,14 +125,14 @@ module Agents
           # Validate that the target agent is in our registry
           # This prevents handoffs to agents that weren't explicitly provided
           unless registry[next_agent.name]
-            puts "[Agents] Warning: Handoff to unregistered agent '#{next_agent.name}', continuing with current agent"
-            # Return the halt content as the final response
             save_conversation_state(chat, context_wrapper, current_agent)
+            error = AgentNotFoundError.new("Handoff failed: Agent '#{next_agent.name}' not found in registry")
             return RunResult.new(
-              output: response.content,
-              messages: MessageExtractor.extract_messages(chat, current_agent),
+              output: nil,
+              messages: Helpers::MessageExtractor.extract_messages(chat, current_agent),
               usage: context_wrapper.usage,
-              context: context_wrapper.context
+              context: context_wrapper.context,
+              error: error
             )
           end
 
@@ -139,9 +146,11 @@ module Agents
           current_agent = next_agent
           context_wrapper.context[:current_agent] = next_agent.name
 
-          # Create new chat for new agent with restored history
-          chat = create_chat(current_agent, context_wrapper)
-          restore_conversation_history(chat, context_wrapper)
+          # Reconfigure existing chat for new agent - preserves conversation history automatically
+          configure_chat_for_agent(chat, current_agent, context_wrapper, replace: true)
+          agent_headers = Helpers::Headers.normalize(current_agent.headers)
+          current_headers = Helpers::Headers.merge(agent_headers, runtime_headers)
+          apply_headers(chat, current_headers)
 
           # Force the new agent to respond to the conversation context
           # This ensures the user gets a response from the new agent
@@ -154,7 +163,7 @@ module Agents
           save_conversation_state(chat, context_wrapper, current_agent)
           return RunResult.new(
             output: response.content,
-            messages: MessageExtractor.extract_messages(chat, current_agent),
+            messages: Helpers::MessageExtractor.extract_messages(chat, current_agent),
             usage: context_wrapper.usage,
             context: context_wrapper.context
           )
@@ -170,7 +179,7 @@ module Agents
 
         return RunResult.new(
           output: response.content,
-          messages: MessageExtractor.extract_messages(chat, current_agent),
+          messages: Helpers::MessageExtractor.extract_messages(chat, current_agent),
           usage: context_wrapper.usage,
           context: context_wrapper.context
         )
@@ -181,7 +190,7 @@ module Agents
 
       RunResult.new(
         output: "Conversation ended: #{e.message}",
-        messages: chat ? MessageExtractor.extract_messages(chat, current_agent) : [],
+        messages: chat ? Helpers::MessageExtractor.extract_messages(chat, current_agent) : [],
         usage: context_wrapper.usage,
         error: e,
         context: context_wrapper.context
@@ -192,7 +201,7 @@ module Agents
 
       RunResult.new(
         output: nil,
-        messages: chat ? MessageExtractor.extract_messages(chat, current_agent) : [],
+        messages: chat ? Helpers::MessageExtractor.extract_messages(chat, current_agent) : [],
         usage: context_wrapper.usage,
         error: e,
         context: context_wrapper.context
@@ -201,6 +210,11 @@ module Agents
 
     private
 
+    # Creates a deep copy of context data for thread safety.
+    # Preserves conversation history array structure while avoiding agent mutation.
+    #
+    # @param context [Hash] The context to copy
+    # @return [Hash] Thread-safe deep copy of the context
     def deep_copy_context(context)
       # Handle deep copying for thread safety
       context.dup.tap do |copied|
@@ -211,13 +225,18 @@ module Agents
       end
     end
 
+    # Restores conversation history from context into RubyLLM chat.
+    # Converts stored message hashes back into RubyLLM::Message objects with proper content handling.
+    #
+    # @param chat [RubyLLM::Chat] The chat instance to restore history into
+    # @param context_wrapper [RunContext] Context containing conversation history
     def restore_conversation_history(chat, context_wrapper)
       history = context_wrapper.context[:conversation_history] || []
 
       history.each do |msg|
         # Only restore user and assistant messages with content
         next unless %i[user assistant].include?(msg[:role].to_sym)
-        next unless msg[:content] && !MessageExtractor.content_empty?(msg[:content])
+        next unless msg[:content] && !Helpers::MessageExtractor.content_empty?(msg[:content])
 
         # Extract text content safely - handle both string and hash content
         content = RubyLLM::Content.new(msg[:content])
@@ -228,21 +247,18 @@ module Agents
           content: content
         )
         chat.add_message(message)
-      rescue StandardError => e
-        # Continue with partial history on error
-        # TODO: Remove this, and let the error propagate up the call stack
-        puts "[Agents] Failed to restore message: #{e.message}\n#{e.backtrace.join("\n")}"
       end
-    rescue StandardError => e
-      # If history restoration completely fails, continue with empty history
-      # TODO: Remove this, and let the error propagate up the call stack
-      puts "[Agents] Failed to restore conversation history: #{e.message}"
-      context_wrapper.context[:conversation_history] = []
     end
 
+    # Saves current conversation state from RubyLLM chat back to context for persistence.
+    # Maintains conversation continuity across agent handoffs and process boundaries.
+    #
+    # @param chat [RubyLLM::Chat] The chat instance to extract state from
+    # @param context_wrapper [RunContext] Context to save state into
+    # @param current_agent [Agents::Agent] The currently active agent
     def save_conversation_state(chat, context_wrapper, current_agent)
       # Extract messages from chat
-      messages = MessageExtractor.extract_messages(chat, current_agent)
+      messages = Helpers::MessageExtractor.extract_messages(chat, current_agent)
 
       # Update context with latest state
       context_wrapper.context[:conversation_history] = messages
@@ -254,14 +270,45 @@ module Agents
       context_wrapper.context.delete(:pending_handoff)
     end
 
-    def create_chat(agent, context_wrapper)
+    # Configures a RubyLLM chat instance with agent-specific settings.
+    # Uses RubyLLM's replace option to swap agent context while preserving conversation history during handoffs.
+    #
+    # @param chat [RubyLLM::Chat] The chat instance to configure
+    # @param agent [Agents::Agent] The agent whose configuration to apply
+    # @param context_wrapper [RunContext] Thread-safe context wrapper
+    # @param replace [Boolean] Whether to replace existing configuration (true for handoffs, false for initial setup)
+    # @return [RubyLLM::Chat] The configured chat instance
+    def configure_chat_for_agent(chat, agent, context_wrapper, replace: false)
       # Get system prompt (may be dynamic)
       system_prompt = agent.get_system_prompt(context_wrapper)
 
-      # Create standard RubyLLM chat
-      chat = RubyLLM::Chat.new(model: agent.model)
-
       # Combine all tools - both handoff and regular tools need wrapping
+      all_tools = build_agent_tools(agent, context_wrapper)
+
+      # Switch model if different (important for handoffs between agents using different models)
+      chat.with_model(agent.model) if replace
+
+      # Configure chat with instructions, temperature, tools, and schema
+      chat.with_instructions(system_prompt, replace: replace) if system_prompt
+      chat.with_temperature(agent.temperature) if agent.temperature
+      chat.with_tools(*all_tools, replace: replace)
+      chat.with_schema(agent.response_schema) if agent.response_schema
+
+      chat
+    end
+
+    def apply_headers(chat, headers)
+      return if headers.empty?
+
+      chat.with_headers(**headers)
+    end
+
+    # Builds thread-safe tool wrappers for an agent's tools and handoff tools.
+    #
+    # @param agent [Agents::Agent] The agent whose tools to wrap
+    # @param context_wrapper [RunContext] Thread-safe context wrapper for tool execution
+    # @return [Array<ToolWrapper>] Array of wrapped tools ready for RubyLLM
+    def build_agent_tools(agent, context_wrapper)
       all_tools = []
 
       # Add handoff tools
@@ -275,13 +322,7 @@ module Agents
         all_tools << ToolWrapper.new(tool, context_wrapper)
       end
 
-      # Configure chat with instructions, temperature, tools, and schema
-      chat.with_instructions(system_prompt) if system_prompt
-      chat.with_temperature(agent.temperature) if agent.temperature
-      chat.with_tools(*all_tools) if all_tools.any?
-      chat.with_schema(agent.response_schema) if agent.response_schema
-
-      chat
+      all_tools
     end
   end
 end

data/lib/agents/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Agents
-  VERSION = "0.5.0"
+  VERSION = "0.6.0"
 end

data/lib/agents.rb

@@ -111,11 +111,11 @@ require_relative "agents/run_context"
 require_relative "agents/tool_context"
 require_relative "agents/tool"
 require_relative "agents/handoff"
+require_relative "agents/helpers"
 require_relative "agents/agent"
 
 # Execution components
 require_relative "agents/tool_wrapper"
-require_relative "agents/message_extractor"
 require_relative "agents/callback_manager"
 require_relative "agents/agent_runner"
 require_relative "agents/runner"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: ai-agents
 version: !ruby/object:Gem::Version
-  version: 0.5.0
+  version: 0.6.0
 platform: ruby
 authors:
 - Shivam Mishra
@@ -15,14 +15,14 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 1.6.0
+        version: 1.8.2
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 1.6.0
+        version: 1.8.2
 description: Ruby AI Agents SDK enables creating complex AI workflows with multi-agent
   orchestration, tool execution, safety guardrails, and provider-agnostic LLM integration.
 email:
@@ -60,6 +60,7 @@ files:
 - docs/guides/agent-as-tool-pattern.md
 - docs/guides/multi-agent-systems.md
 - docs/guides/rails-integration.md
+- docs/guides/request-headers.md
 - docs/guides/state-persistence.md
 - docs/guides/structured-output.md
 - docs/index.md
@@ -102,7 +103,9 @@ files:
 - lib/agents/agent_tool.rb
 - lib/agents/callback_manager.rb
 - lib/agents/handoff.rb
-- lib/agents/message_extractor.rb
+- lib/agents/helpers.rb
+- lib/agents/helpers/headers.rb
+- lib/agents/helpers/message_extractor.rb
 - lib/agents/result.rb
 - lib/agents/run_context.rb
 - lib/agents/runner.rb

data/lib/agents/message_extractor.rb

@@ -1,97 +0,0 @@
-# frozen_string_literal: true
-
-module Agents
-  # Service object responsible for extracting and formatting conversation messages
-  # from RubyLLM chat objects into a format suitable for persistence and context restoration.
-  #
-  # Handles different message types:
-  # - User messages: Basic content preservation
-  # - Assistant messages: Includes agent attribution and tool calls
-  # - Tool result messages: Links back to original tool calls
-  #
-  # @example Extract messages from a chat
-  #   messages = MessageExtractor.extract_messages(chat, current_agent)
-  #   #=> [
-  #     { role: :user, content: "Hello" },
-  #     { role: :assistant, content: "Hi!", agent_name: "Support", tool_calls: [...] },
-  #     { role: :tool, content: "Result", tool_call_id: "call_123" }
-  #   ]
-  class MessageExtractor
-    # Check if content is considered empty (handles both String and Hash content)
-    #
-    # @param content [String, Hash, nil] The content to check
-    # @return [Boolean] true if content is empty, false otherwise
-    def self.content_empty?(content)
-      case content
-      when String
-        content.strip.empty?
-      when Hash
-        content.empty?
-      else
-        content.nil?
-      end
-    end
-
-    # Extract messages from a chat object for conversation history persistence
-    #
-    # @param chat [Object] Chat object that responds to :messages
-    # @param current_agent [Agent] The agent currently handling the conversation
-    # @return [Array<Hash>] Array of message hashes suitable for persistence
-    def self.extract_messages(chat, current_agent)
-      new(chat, current_agent).extract
-    end
-
-    def initialize(chat, current_agent)
-      @chat = chat
-      @current_agent = current_agent
-    end
-
-    def extract
-      return [] unless @chat.respond_to?(:messages)
-
-      @chat.messages.filter_map do |msg|
-        case msg.role
-        when :user, :assistant
-          extract_user_or_assistant_message(msg)
-        when :tool
-          extract_tool_message(msg)
-        end
-      end
-    end
-
-    private
-
-    def extract_user_or_assistant_message(msg)
-      return nil unless msg.content && !self.class.content_empty?(msg.content)
-
-      message = {
-        role: msg.role,
-        content: msg.content
-      }
-
-      if msg.role == :assistant
-        # Add agent attribution for conversation continuity
-        message[:agent_name] = @current_agent.name if @current_agent
-
-        # Add tool calls if present
-        if msg.tool_call? && msg.tool_calls
-          # RubyLLM stores tool_calls as Hash with call_id => ToolCall object
-          # Reference: RubyLLM::StreamAccumulator#tool_calls_from_stream
-          message[:tool_calls] = msg.tool_calls.values.map(&:to_h)
-        end
-      end
-
-      message
-    end
-
-    def extract_tool_message(msg)
-      return nil unless msg.tool_result?
-
-      {
-        role: msg.role,
-        content: msg.content,
-        tool_call_id: msg.tool_call_id
-      }
-    end
-  end
-end