Fix callback persistence chaining using Fanout pattern

lib/ruby_llm/active_record/acts_as_legacy.rb

@@ -152,26 +152,12 @@ def with_schema(...)
       end
 
       def on_new_message(&block)
-        to_llm
-
-        existing_callback = @chat.instance_variable_get(:@on)[:new_message]
-
-        @chat.on_new_message do
-          existing_callback&.call
-          block&.call
-        end
+        to_llm.on_new_message(&block)
         self
       end
 
       def on_end_message(&block)
-        to_llm
-
-        existing_callback = @chat.instance_variable_get(:@on)[:end_message]
-
-        @chat.on_end_message do |msg|
-          existing_callback&.call(msg)
-          block&.call(msg)
-        end
+        to_llm.on_end_message(&block)
         self
       end
 

lib/ruby_llm/active_record/chat_methods.rb

@@ -140,26 +140,12 @@ def with_schema(...)
       end
 
       def on_new_message(&block)
-        to_llm
-
-        existing_callback = @chat.instance_variable_get(:@on)[:new_message]
-
-        @chat.on_new_message do
-          existing_callback&.call
-          block&.call
-        end
+        to_llm.on_new_message(&block)
         self
       end
 
       def on_end_message(&block)
-        to_llm
-
-        existing_callback = @chat.instance_variable_get(:@on)[:end_message]
-
-        @chat.on_end_message do |msg|
-          existing_callback&.call(msg)
-          block&.call(msg)
-        end
+        to_llm.on_end_message(&block)
         self
       end
 

lib/ruby_llm/chat.rb

@@ -7,6 +7,26 @@ class Chat
 
     attr_reader :model, :messages, :tools, :params, :headers, :schema
 
+    # Stores multiple callbacks per key and invokes all of them
+    class CallbackFanout
+      def initialize
+        @callbacks = {}
+      end
+
+      def [](key)
+        callbacks = @callbacks[key]
+        return if callbacks.nil? || callbacks.empty?
+
+        ->(*args) { callbacks.each { |cb| cb.call(*args) } }
+      end
+
+      def []=(key, callable)
+        return unless callable
+
+        (@callbacks[key] ||= []) << callable
+      end
+    end
+
     def initialize(model: nil, provider: nil, assume_model_exists: false, context: nil)
       if assume_model_exists && !provider
         raise ArgumentError, 'Provider must be specified if assume_model_exists is true'
@@ -22,12 +42,7 @@ def initialize(model: nil, provider: nil, assume_model_exists: false, context: n
       @params = {}
       @headers = {}
       @schema = nil
-      @on = {
-        new_message: nil,
-        end_message: nil,
-        tool_call: nil,
-        tool_result: nil
-      }
+      @on = CallbackFanout.new
     end
 
     def ask(message = nil, with: nil, &)

spec/ruby_llm/active_record/acts_as_spec.rb

@@ -622,6 +622,20 @@ def uploaded_file(path, type)
       expect(chat.messages.count).to eq(2) # Persistence still works
     end
 
+    it 'allows chaining callbacks on to_llm without losing persistence' do
+      chat = Chat.create!(model: model)
+      llm_chat = chat.to_llm
+
+      user_callback_called = false
+      # Directly attach callback to the underlying Chat object
+      llm_chat.on_new_message { user_callback_called = true }
+
+      chat.ask('Hello')
+
+      expect(user_callback_called).to be true
+      expect(chat.messages.count).to eq(2) # Persistence still works
+    end
+
     it 'calls on_tool_call and on_tool_result callbacks' do
       tool_call_received = nil
       tool_result_received = nil