feat: replace based chat configuration (#33)

Author: scmmishra

.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

lib/agents/runner.rb

@@ -55,6 +55,7 @@ class Runner
     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.
@@ -91,14 +92,15 @@ def run(starting_agent, input, context: {}, registry: {}, max_turns: DEFAULT_MAX
       current_turn = 0
 
       # Create chat and restore conversation history
-      chat = create_chat(current_agent, context_wrapper)
+      chat = RubyLLM::Chat.new(model: current_agent.model)
+      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 +120,14 @@ def run(starting_agent, input, context: {}, registry: {}, max_turns: DEFAULT_MAX
           # 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,
+              output: nil,
               messages: MessageExtractor.extract_messages(chat, current_agent),
               usage: context_wrapper.usage,
-              context: context_wrapper.context
+              context: context_wrapper.context,
+              error: error
             )
           end
 
@@ -139,9 +141,8 @@ def run(starting_agent, input, context: {}, registry: {}, max_turns: DEFAULT_MAX
           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)
 
           # Force the new agent to respond to the conversation context
           # This ensures the user gets a response from the new agent
@@ -201,6 +202,11 @@ def run(starting_agent, input, context: {}, registry: {}, max_turns: DEFAULT_MAX
 
     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,6 +217,11 @@ def deep_copy_context(context)
       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] || []
 
@@ -228,18 +239,15 @@ def restore_conversation_history(chat, context_wrapper)
           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)
@@ -254,14 +262,39 @@ def save_conversation_state(chat, context_wrapper, current_agent)
       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
+
+    # 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 +308,7 @@ def create_chat(agent, context_wrapper)
         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

spec/agents/runner_spec.rb

@@ -196,6 +196,50 @@
         expect(result.output).to eq("Hello, I'm the specialist. How can I help?")
         expect(result.context[:current_agent]).to eq("HandoffAgent")
       end
+
+      it "returns error when handoff to unregistered agent is attempted" do
+        # Only register the triage agent, not the handoff target
+        registry = { "TriageAgent" => agent_with_handoffs }
+
+        # Mock only the first tool call that triggers handoff
+        stub_request(:post, "https://api.openai.com/v1/chat/completions")
+          .to_return(
+            status: 200,
+            body: {
+              id: "chatcmpl-handoff",
+              object: "chat.completion",
+              created: 1_677_652_288,
+              model: "gpt-4o",
+              choices: [{
+                index: 0,
+                message: {
+                  role: "assistant",
+                  content: nil,
+                  tool_calls: [{
+                    id: "call_handoff",
+                    type: "function",
+                    function: {
+                      name: "handoff_to_handoffagent",
+                      arguments: "{}"
+                    }
+                  }]
+                },
+                finish_reason: "tool_calls"
+              }],
+              usage: { prompt_tokens: 20, completion_tokens: 5, total_tokens: 25 }
+            }.to_json,
+            headers: { "Content-Type" => "application/json" }
+          )
+
+        result = runner.run(agent_with_handoffs, "I need specialist help", registry: registry)
+
+        expect(result.failed?).to be true
+        expect(result.error).to be_a(Agents::Runner::AgentNotFoundError)
+        expect(result.error.message).to eq("Handoff failed: Agent 'HandoffAgent' not found in registry")
+        expect(result.output).to be_nil
+        expect(result.context[:current_agent]).to eq("TriageAgent")
+        expect(result.context[:pending_handoff]).to be_nil # Should clear pending handoff
+      end
     end
 
     context "when max_turns is exceeded" do
@@ -204,8 +248,9 @@
         mock_chat = instance_double(RubyLLM::Chat)
         mock_response = instance_double(RubyLLM::Message, tool_call?: true)
 
+        allow(RubyLLM::Chat).to receive(:new).and_return(mock_chat)
         allow(runner).to receive_messages(
-          create_chat: mock_chat,
+          configure_chat_for_agent: mock_chat,
           restore_conversation_history: nil,
           save_conversation_state: nil
         )
@@ -224,7 +269,7 @@
     context "when standard error occurs" do
       it "handles errors gracefully and returns error result" do
         # Mock chat creation to raise an error
-        allow(runner).to receive(:create_chat).and_raise(StandardError, "Test error")
+        allow(RubyLLM::Chat).to receive(:new).and_raise(StandardError, "Test error")
 
         result = runner.run(agent, "Error test")
 
@@ -272,8 +317,9 @@
         mock_halt = instance_double(RubyLLM::Tool::Halt, content: "Processing complete", is_a?: true)
 
         allow(mock_halt).to receive(:is_a?).with(RubyLLM::Tool::Halt).and_return(true)
+        allow(RubyLLM::Chat).to receive(:new).and_return(mock_chat)
         allow(runner).to receive_messages(
-          create_chat: mock_chat,
+          configure_chat_for_agent: mock_chat,
           restore_conversation_history: nil,
           save_conversation_state: nil
         )
@@ -368,5 +414,37 @@
         end
       end
     end
+
+    context "when agent has regular tools" do
+      let(:agent_with_tools) do
+        instance_double(Agents::Agent,
+                        name: "ToolAgent",
+                        model: "gpt-4o",
+                        tools: [test_tool],
+                        handoff_agents: [],
+                        temperature: 0.7,
+                        response_schema: nil,
+                        get_system_prompt: "You are an agent with tools")
+      end
+
+      it "wraps regular tools in ToolWrapper" do
+        # Spy on ToolWrapper constructor
+        allow(Agents::ToolWrapper).to receive(:new).and_call_original
+
+        # Stub a simple response that doesn't use tools
+        stub_simple_chat("I have tools available")
+
+        runner.run(
+          agent_with_tools,
+          "Hello",
+          context: {},
+          registry: { "ToolAgent" => agent_with_tools },
+          max_turns: 1
+        )
+
+        # Verify ToolWrapper was called with the regular tool
+        expect(Agents::ToolWrapper).to have_received(:new).with(test_tool, anything)
+      end
+    end
   end
 end

spec/agents/tool_context_spec.rb

@@ -48,4 +48,51 @@
       expect(tool_context.retry_count).to eq(2)
     end
   end
+
+  describe "#state" do
+    context "when state exists in context" do
+      let(:existing_state) { { customer_id: 456, customer_name: "John" } }
+      let(:context_hash) { { user_id: 123, session: "test", state: existing_state } }
+
+      it "returns the existing state hash" do
+        result = tool_context.state
+        expect(result).to eq(existing_state)
+      end
+
+      it "allows modifications to the state" do
+        tool_context.state[:new_key] = "new_value"
+        expect(tool_context.state[:new_key]).to eq("new_value")
+      end
+    end
+
+    context "when state does not exist in context" do
+      let(:context_hash) { { user_id: 123, session: "test" } }
+
+      it "initializes state as empty hash" do
+        result = tool_context.state
+        expect(result).to eq({})
+      end
+
+      it "persists the initialized state in context" do
+        tool_context.state
+        expect(context_hash[:state]).to eq({})
+      end
+
+      it "allows adding to the initialized state" do
+        tool_context.state[:customer_id] = 789
+        expect(tool_context.state[:customer_id]).to eq(789)
+        expect(context_hash[:state][:customer_id]).to eq(789)
+      end
+    end
+
+    context "when state is nil in context" do
+      let(:context_hash) { { user_id: 123, session: "test", state: nil } }
+
+      it "replaces nil with empty hash" do
+        result = tool_context.state
+        expect(result).to eq({})
+        expect(context_hash[:state]).to eq({})
+      end
+    end
+  end
 end