context: always emit on_tool_return

Before this change: in order for `on_tool_return` to be triggered,
you had to schedule a tool via `LLM::Stream#queue`.

After this change: all tool calls go through `on_tool_return`,
regardless of how they were scheduled.

Author: bsdrobert

mrblib/mruby-llm/agent.rb

@@ -437,6 +437,7 @@ def call_functions
       results = confirmables.map do |tool|
         send(:on_tool_confirmation, tool, strategy)
       end
+      @ctx.method(:emit_tool_returns).call(confirmables, results)
       @ctx.functions? ? [*results, *wait(strategy)] : results
     end
 

mrblib/mruby-llm/context.rb

@@ -272,8 +272,11 @@ def wait(strategy)
         @queue.wait
       else
         return guarded_returns if guarded_returns
-        @queue = functions.spawn(strategy)
-        @queue.wait
+        tools = functions
+        @queue = tools.spawn(strategy)
+        returns = @queue.wait
+        emit_tool_returns(tools, returns)
+        returns
       end
     ensure
       @queue = nil
@@ -529,6 +532,14 @@ def guarded_return_for(function, warning)
       })
     end
 
+    ##
+    # Emits tool return callbacks for directly waited function work.
+    # @api private
+    def emit_tool_returns(tools, returns)
+      return unless LLM::Stream === stream
+      returns.each_with_index { |result, index| stream.on_tool_return(tools[index], result) }
+    end
+
     ##
     # Closes assistant tool-call messages that do not have matching tool
     # responses. This can happen when a turn is interrupted while a tool call is

mrblib/mruby-llm/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module LLM
-  VERSION = "0.1.0.beta.7"
+  VERSION = "0.1.0.beta.8"
 end

spec/agent_spec.rb

@@ -248,14 +248,6 @@ def configured_model
   end
 
   describe "#talk" do
-    before do
-      transport.stub(
-        "POST", "/v1/chat/completions",
-        fixture: "openai/chat_completions.sse",
-        headers: {"content-type" => "text/event-stream"}
-      )
-    end
-
     context "when configured with a class-level stream object" do
       let(:stream) { StringIO.new }
       let(:agent_class) do
@@ -267,6 +259,11 @@ def configured_model
       let(:agent) { agent_class.new(llm, model: "gpt-4.1") }
 
       before do
+        transport.stub(
+          "POST", "/v1/chat/completions",
+          fixture: "openai/chat_completions.sse",
+          headers: {"content-type" => "text/event-stream"}
+        )
         agent.talk("Say hello")
       end
 
@@ -287,6 +284,11 @@ def configured_model
       let(:agent) { agent_class.new(llm, model: "gpt-4.1", stream: override_stream) }
 
       before do
+        transport.stub(
+          "POST", "/v1/chat/completions",
+          fixture: "openai/chat_completions.sse",
+          headers: {"content-type" => "text/event-stream"}
+        )
         agent.talk("Say hello")
       end
 
@@ -298,6 +300,102 @@ def configured_model
         expect(default_stream.string).must_equal ""
       end
     end
+
+    context "when auto-executing streamed tool calls" do
+      let(:stream) do
+        Class.new(LLM::Stream) do
+          attr_reader :returns
+          def initialize
+            @returns = []
+          end
+          def on_tool_return(tool, result)
+            @returns << [tool.name, result.name, result.value]
+          end
+        end.new
+      end
+      let(:tool) do
+        Class.new(LLM::Tool) do
+          name "system"
+          parameter :command, String, "The command"
+          required %i[command]
+          def call(command:)
+            {"success" => command == "date" ? "2025-08-24" : false}
+          end
+        end
+      end
+      let(:agent) { LLM::Agent.new(llm, model: "gpt-4.1", stream:, tools: [tool]) }
+
+      before do
+        transport
+          .stub(
+            "POST", "/v1/chat/completions",
+            fixture: "openai/chat_completions_tool.sse",
+            headers: {"content-type" => "text/event-stream"}
+          )
+          .stub(
+            "POST", "/v1/chat/completions",
+            fixture: "openai/chat_completions_tool_result.sse",
+            headers: {"content-type" => "text/event-stream"}
+          )
+        agent.talk("Run date")
+      end
+
+      it "emits tool returns without manual stream queueing" do
+        expect(stream.returns).must_equal [["system", "system", {"success" => "2025-08-24"}]]
+      end
+    end
+
+    context "when confirming streamed tool calls" do
+      let(:stream) do
+        Class.new(LLM::Stream) do
+          attr_reader :returns
+          def initialize
+            @returns = []
+          end
+          def on_tool_return(tool, result)
+            @returns << [tool.name, result.name, result.value]
+          end
+        end.new
+      end
+      let(:tool) do
+        Class.new(LLM::Tool) do
+          name "system"
+          parameter :command, String, "The command"
+          required %i[command]
+          def call(command:)
+            {"success" => command == "date" ? "2025-08-24" : false}
+          end
+        end
+      end
+      let(:agent_class) do
+        Class.new(LLM::Agent) do
+          confirm "system"
+          def on_tool_confirmation(fn, strategy)
+            fn.spawn(strategy).wait
+          end
+        end
+      end
+      let(:agent) { agent_class.new(llm, model: "gpt-4.1", stream:, tools: [tool]) }
+
+      before do
+        transport
+          .stub(
+            "POST", "/v1/chat/completions",
+            fixture: "openai/chat_completions_tool.sse",
+            headers: {"content-type" => "text/event-stream"}
+          )
+          .stub(
+            "POST", "/v1/chat/completions",
+            fixture: "openai/chat_completions_tool_result.sse",
+            headers: {"content-type" => "text/event-stream"}
+          )
+        agent.talk("Run date")
+      end
+
+      it "emits confirmed tool returns" do
+        expect(stream.returns).must_equal [["system", "system", {"success" => "2025-08-24"}]]
+      end
+    end
   end
 end
 

spec/context_spec.rb

@@ -145,6 +145,17 @@ def call(command:)
   end
 
   context "when given a streamed tool call" do
+    let(:stream) do
+      Class.new(LLM::Stream) do
+        attr_reader :returns
+        def initialize
+          @returns = []
+        end
+        def on_tool_return(tool, result)
+          @returns << [tool.name, result.name, result.value]
+        end
+      end.new
+    end
     let(:tool) do
       Class.new(LLM::Tool) do
         name "system"
@@ -157,7 +168,7 @@ def call(command:)
         end
       end
     end
-    let(:params) { {model: "gpt-4.1", stream: true, tools: [tool]} }
+    let(:params) { {model: "gpt-4.1", stream:, tools: [tool]} }
     let(:prompt) do
       ctx.build_prompt do
         _1.talk "You are a bot that can run UNIX system commands"
@@ -186,6 +197,11 @@ def call(command:)
       expect(res.content).must_equal "Today's date is August 24, 2025."
       expect(ctx.messages.last.content).must_equal "Today's date is August 24, 2025."
     end
+
+    it "emits tool returns from direct waits" do
+      ctx.wait(:call)
+      expect(stream.returns).must_equal [["system", "system", {"success" => "2025-08-24"}]]
+    end
   end
 
   context "when resuming a serialized conversation" do