DUX-10302: Fix Bedrock tool-call argument routing by contentBlockIndex

lib/ruby_llm/error.rb

@@ -39,6 +39,28 @@ def initialize(type = nil)
     end
   end
 
+  # Raised when a streamed tool call's accumulated arguments are not valid JSON, e.g.
+  # because a delta fragment was routed to the wrong content block or the stream ended
+  # mid-argument. `finish_reason` lets callers tell a genuine output-token-cap truncation
+  # (finish_reason indicates max tokens) apart from a "complete" turn whose tool-call
+  # arguments are corrupt regardless (a routing bug or wire error, not fixable by a bigger
+  # token budget).
+  class ToolCallArgumentsTruncatedError < StandardError
+    attr_reader :tool_call_id, :tool_name, :raw_arguments, :finish_reason
+
+    def initialize(tool_call_id:, tool_name:, raw_arguments:, finish_reason:)
+      @tool_call_id = tool_call_id
+      @tool_name = tool_name
+      @raw_arguments = raw_arguments
+      @finish_reason = finish_reason
+
+      super(
+        "Tool call #{tool_name.inspect} (id: #{tool_call_id.inspect}) has incomplete or unparseable " \
+        "arguments (finish_reason: #{finish_reason.inspect}): #{raw_arguments.inspect}"
+      )
+    end
+  end
+
   # Error classes for different HTTP status codes
   # Raised when the API request is invalid.
   class BadRequestError < Error

lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb

@@ -180,9 +180,15 @@ def build_content_block_start_chunk(event, thinking_state)
 
           case content_block['type']
           when 'tool_use'
+            # Keyed by index (mirrors Converse::Streaming#extract_tool_call_start) so the
+            # accumulator's index-routing can tell parallel tool-use blocks apart — without it,
+            # every delta falls back to whichever tool call started most recently, corrupting
+            # arguments whenever a later block's content_block_start arrives before an earlier
+            # block's remaining deltas. Falls back to the tool_use id if index is ever absent.
             id = content_block['id']
+            stream_key = index.nil? ? id : index
             tool_calls = {
-              id => ToolCall.new(id: id, name: content_block['name'], arguments: {})
+              stream_key => ToolCall.new(id: id, name: content_block['name'], arguments: {})
             }
           when 'redacted_thinking'
             thinking = Thinking.build(blocks: [{ 'type' => 'redacted_thinking', 'data' => content_block['data'] }])
@@ -213,8 +219,10 @@ def build_content_block_delta_chunk(event, thinking_state)
           when 'text_delta'
             content = delta['text']
           when 'input_json_delta'
+            # See build_content_block_start_chunk. Falls back to nil (today's "latest started
+            # tool call" behavior) on the rare event with no index.
             partial = delta['partial_json']
-            tool_calls = { nil => ToolCall.new(id: nil, name: nil, arguments: partial) } if partial
+            tool_calls = { index => ToolCall.new(id: nil, name: nil, arguments: partial) } if partial
           when 'thinking_delta'
             thinking_text = delta['thinking']
             thinking_state[index][:text] << thinking_text.to_s if thinking_state[index]

lib/ruby_llm/protocols/converse/streaming.rb

@@ -384,25 +384,36 @@ def tool_call_delta_event?(event)
           event['contentBlockDelta'] || event.dig('delta', 'toolUse')
         end
 
+        # Keyed by contentBlockIndex when present so the accumulator's index-routing
+        # (@tool_call_ids_by_index) can tell parallel tool-use blocks apart — without it,
+        # every delta falls back to whichever tool call started most recently, corrupting
+        # arguments whenever a later block's contentBlockStart arrives before an earlier
+        # block's remaining deltas. Falls back to the toolUseId (today's behavior) for the
+        # legacy/test-only flat 'start' shape, which carries no contentBlockIndex.
         def extract_tool_call_start(event)
           tool_use = event.dig('contentBlockStart', 'start', 'toolUse') || event.dig('start', 'toolUse')
           return nil unless tool_use
 
           tool_use_id = tool_use['toolUseId']
+          stream_key = event.dig('contentBlockStart', 'contentBlockIndex') || tool_use_id
           {
-            tool_use_id => ToolCall.new(
+            stream_key => ToolCall.new(
               id: tool_use_id,
               name: tool_use['name'],
               arguments: tool_use['input'] || {}
             )
           }
         end
 
+        # See extract_tool_call_start. Falls back to nil (today's "latest started tool
+        # call" behavior) for the legacy/test-only flat 'delta' shape, which carries no
+        # contentBlockIndex.
         def extract_tool_call_delta(event)
           input = normalized_delta(event).dig('toolUse', 'input')
           return nil unless input
 
-          { nil => ToolCall.new(id: nil, name: nil, arguments: input) }
+          stream_key = event.dig('contentBlockDelta', 'contentBlockIndex')
+          { stream_key => ToolCall.new(id: nil, name: nil, arguments: input) }
         end
 
         def normalized_delta(event)

lib/ruby_llm/stream_accumulator.rb

@@ -86,7 +86,7 @@ def resolve_citation_text(citation)
     def tool_calls_from_stream
       tool_calls.transform_values do |tc|
         arguments = if tc.arguments.is_a?(String) && !tc.arguments.empty?
-                      JSON.parse(tc.arguments)
+                      parse_tool_call_arguments(tc)
                     elsif tc.arguments.is_a?(String)
                       {}
                     else
@@ -102,6 +102,24 @@ def tool_calls_from_stream
       end
     end
 
+    # A tool call's accumulated argument fragments should always join into valid JSON;
+    # a parse failure here means fragments were dropped or misrouted (e.g. a content-block
+    # index bug) or the provider ended the stream mid-argument. Raising — rather than
+    # silently executing a tool with half-parsed arguments, or silently swallowing to `{}`
+    # — surfaces exactly which tool call is affected. @finish_reason lets the caller
+    # distinguish this from a legitimate output-token-cap truncation (see
+    # ToolCallArgumentsTruncatedError).
+    def parse_tool_call_arguments(tool_call)
+      JSON.parse(tool_call.arguments)
+    rescue JSON::ParserError
+      raise ToolCallArgumentsTruncatedError.new(
+        tool_call_id: tool_call.id,
+        tool_name: tool_call.name,
+        raw_arguments: tool_call.arguments,
+        finish_reason: @finish_reason
+      )
+    end
+
     def accumulate_tool_calls(new_tool_calls)
       RubyLLM.logger.debug { "Accumulating tool calls: #{new_tool_calls}" } if RubyLLM.config.log_stream_debug
       new_tool_calls.each do |stream_key, tool_call|
@@ -136,7 +154,12 @@ def initial_tool_call_arguments(tool_call)
 
     def append_tool_call_fragment(stream_key, tool_call)
       existing = find_tool_call(stream_key)
-      return unless existing
+      unless existing
+        if RubyLLM.config.log_stream_debug
+          RubyLLM.logger.debug { "Dropping tool call fragment for unresolved stream_key: #{stream_key.inspect}" }
+        end
+        return
+      end
 
       fragment = tool_call.arguments
       fragment = '' if fragment.nil?

spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb

@@ -518,7 +518,7 @@ def render_payload(messages = [], **opts)
       }
       chunk = streaming.send(:build_chunk, event)
       expect(chunk.tool_calls).not_to be_nil
-      expect(chunk.tool_calls[nil].arguments).to eq('{"key":')
+      expect(chunk.tool_calls[1].arguments).to eq('{"key":')
     end
 
     it 'extracts thinking text from thinking_delta' do
@@ -569,7 +569,8 @@ def render_payload(messages = [], **opts)
       }
       chunk = streaming.send(:build_chunk, event)
       expect(chunk.tool_calls).not_to be_nil
-      tc = chunk.tool_calls['call_xyz']
+      tc = chunk.tool_calls[0]
+      expect(tc.id).to eq('call_xyz')
       expect(tc.name).to eq('search')
     end
 
@@ -594,6 +595,73 @@ def render_payload(messages = [], **opts)
     end
   end
 
+  # ---------------------------------------------------------------------------
+  # Streaming parallel tool calls
+  # ---------------------------------------------------------------------------
+
+  describe 'Streaming parallel tool calls' do
+    let(:streaming) do
+      described_class.allocate.tap do |obj|
+        obj.instance_variable_set(:@model, build_model('anthropic.claude-haiku-4-5-20251001-v1:0'))
+        obj.instance_variable_set(:@config, build_config)
+      end
+    end
+
+    # Feeds events through build_chunk into a single accumulator, mirroring how
+    # stream_response drives the real event stream.
+    def accumulate(events)
+      accumulator = RubyLLM::StreamAccumulator.new
+      thinking_state = {}
+      events.each { |e| accumulator.add(streaming.send(:build_chunk, e, thinking_state)) }
+      accumulator.to_message(nil)
+    end
+
+    it 'keeps a block\'s arguments intact when a delta for it arrives after the next block starts' do
+      events = [
+        { 'type' => 'content_block_start', 'index' => 0,
+          'content_block' => { 'type' => 'tool_use', 'id' => 'call_1', 'name' => 'market_data' } },
+        { 'type' => 'content_block_delta', 'index' => 0,
+          'delta' => { 'type' => 'input_json_delta', 'partial_json' => '{"symbol":"MNQM26",' } },
+        { 'type' => 'content_block_start', 'index' => 1,
+          'content_block' => { 'type' => 'tool_use', 'id' => 'call_2', 'name' => 'search' } },
+        { 'type' => 'content_block_delta', 'index' => 1,
+          'delta' => { 'type' => 'input_json_delta', 'partial_json' => '{"query":"market news"}' } },
+        # A delta for block 0 arriving after block 1 has already started — this is the
+        # exact ordering that corrupted arguments before index was threaded through as
+        # the stream key.
+        { 'type' => 'content_block_delta', 'index' => 0,
+          'delta' => { 'type' => 'input_json_delta', 'partial_json' => '"interval":"minute"}' } },
+        { 'type' => 'content_block_stop', 'index' => 0 },
+        { 'type' => 'content_block_stop', 'index' => 1 }
+      ]
+
+      message = accumulate(events)
+
+      expect(message.tool_calls['call_1'].arguments).to eq('symbol' => 'MNQM26', 'interval' => 'minute')
+      expect(message.tool_calls['call_2'].arguments).to eq('query' => 'market news')
+    end
+
+    it 'accumulates sequential (non-interleaved) parallel tool calls correctly' do
+      events = [
+        { 'type' => 'content_block_start', 'index' => 0,
+          'content_block' => { 'type' => 'tool_use', 'id' => 'call_1', 'name' => 'market_data' } },
+        { 'type' => 'content_block_delta', 'index' => 0,
+          'delta' => { 'type' => 'input_json_delta', 'partial_json' => '{"symbol":"MNQM26"}' } },
+        { 'type' => 'content_block_stop', 'index' => 0 },
+        { 'type' => 'content_block_start', 'index' => 1,
+          'content_block' => { 'type' => 'tool_use', 'id' => 'call_2', 'name' => 'search' } },
+        { 'type' => 'content_block_delta', 'index' => 1,
+          'delta' => { 'type' => 'input_json_delta', 'partial_json' => '{"query":"market news"}' } },
+        { 'type' => 'content_block_stop', 'index' => 1 }
+      ]
+
+      message = accumulate(events)
+
+      expect(message.tool_calls['call_1'].arguments).to eq('symbol' => 'MNQM26')
+      expect(message.tool_calls['call_2'].arguments).to eq('query' => 'market news')
+    end
+  end
+
   # ---------------------------------------------------------------------------
   # Streaming multi-block thinking parse + replay
   # ---------------------------------------------------------------------------

spec/ruby_llm/protocols/converse/streaming_spec.rb

@@ -268,6 +268,89 @@ def accumulate(events)
     end
   end
 
+  describe 'parallel tool calls' do
+    # Feeds events through build_chunk into a single accumulator, mirroring how
+    # stream_response drives the real event stream.
+    def accumulate(events)
+      accumulator = RubyLLM::StreamAccumulator.new
+      events.each { |e| accumulator.add(streaming.send(:build_chunk, e)) }
+      accumulator.to_message(nil)
+    end
+
+    it 'keeps a block\'s arguments intact when a delta for it arrives after the next block starts' do
+      events = [
+        { 'contentBlockStart' => { 'contentBlockIndex' => 0,
+                                   'start' => { 'toolUse' => { 'toolUseId' => 'call_1', 'name' => 'market_data' } } } },
+        { 'contentBlockDelta' => { 'contentBlockIndex' => 0,
+                                   'delta' => { 'toolUse' => { 'input' => '{"symbol":"MNQM26",' } } } },
+        { 'contentBlockStart' => { 'contentBlockIndex' => 1,
+                                   'start' => { 'toolUse' => { 'toolUseId' => 'call_2', 'name' => 'search' } } } },
+        { 'contentBlockDelta' => { 'contentBlockIndex' => 1,
+                                   'delta' => { 'toolUse' => { 'input' => '{"query":"market news"}' } } } },
+        # A delta for block 0 arriving after block 1 has already started — this is the
+        # exact ordering that corrupted arguments before contentBlockIndex was threaded
+        # through as the stream key.
+        { 'contentBlockDelta' => { 'contentBlockIndex' => 0,
+                                   'delta' => { 'toolUse' => { 'input' => '"interval":"minute"}' } } } },
+        { 'contentBlockStop' => { 'contentBlockIndex' => 0 } },
+        { 'contentBlockStop' => { 'contentBlockIndex' => 1 } }
+      ]
+
+      message = accumulate(events)
+
+      expect(message.tool_calls['call_1'].arguments).to eq('symbol' => 'MNQM26', 'interval' => 'minute')
+      expect(message.tool_calls['call_2'].arguments).to eq('query' => 'market news')
+    end
+
+    it 'accumulates sequential (non-interleaved) parallel tool calls correctly' do
+      events = [
+        { 'contentBlockStart' => { 'contentBlockIndex' => 0,
+                                   'start' => { 'toolUse' => { 'toolUseId' => 'call_1', 'name' => 'market_data' } } } },
+        { 'contentBlockDelta' => { 'contentBlockIndex' => 0,
+                                   'delta' => { 'toolUse' => { 'input' => '{"symbol":"MNQM26"}' } } } },
+        { 'contentBlockStop' => { 'contentBlockIndex' => 0 } },
+        { 'contentBlockStart' => { 'contentBlockIndex' => 1,
+                                   'start' => { 'toolUse' => { 'toolUseId' => 'call_2', 'name' => 'search' } } } },
+        { 'contentBlockDelta' => { 'contentBlockIndex' => 1,
+                                   'delta' => { 'toolUse' => { 'input' => '{"query":"market news"}' } } } },
+        { 'contentBlockStop' => { 'contentBlockIndex' => 1 } }
+      ]
+
+      message = accumulate(events)
+
+      expect(message.tool_calls['call_1'].arguments).to eq('symbol' => 'MNQM26')
+      expect(message.tool_calls['call_2'].arguments).to eq('query' => 'market news')
+    end
+
+    it 'still accumulates via the legacy flat start/delta shapes with no contentBlockIndex' do
+      events = [
+        { 'start' => { 'toolUse' => { 'toolUseId' => 'call_1', 'name' => 'weather' } } },
+        { 'delta' => { 'toolUse' => { 'input' => '{"city":"SF"}' } } }
+      ]
+
+      message = accumulate(events)
+
+      expect(message.tool_calls['call_1'].name).to eq('weather')
+      expect(message.tool_calls['call_1'].arguments).to eq('city' => 'SF')
+    end
+
+    it 'raises ToolCallArgumentsTruncatedError with the finish_reason when arguments never close' do
+      events = [
+        { 'contentBlockStart' => { 'contentBlockIndex' => 0,
+                                   'start' => { 'toolUse' => { 'toolUseId' => 'call_1', 'name' => 'market_data' } } } },
+        { 'contentBlockDelta' => { 'contentBlockIndex' => 0,
+                                   'delta' => { 'toolUse' => { 'input' => '{"symbol":"MNQM26"' } } } },
+        { 'messageStop' => { 'stopReason' => 'tool_use' } }
+      ]
+
+      expect { accumulate(events) }.to raise_error(RubyLLM::ToolCallArgumentsTruncatedError) do |error|
+        expect(error.tool_call_id).to eq('call_1')
+        expect(error.tool_name).to eq('market_data')
+        expect(error.finish_reason).to eq('tool_use')
+      end
+    end
+  end
+
   # ConverseStream's real wire format carries the event type in the eventstream
   # :event-type HEADER; the JSON payload is the bare member struct — e.g.
   # {"contentBlockIndex":0,"delta":{...}}, {"stopReason":"tool_use"} — never nested
@@ -357,6 +440,30 @@ def stream_to_message(*typed_payloads)
       )
     end
 
+    it 'keeps parallel tool call arguments intact when flat-framed and interleaved across blocks' do
+      message = stream_to_message(
+        ['contentBlockStart', { 'contentBlockIndex' => 0,
+                                'start' => { 'toolUse' => { 'toolUseId' => 'call_1', 'name' => 'market_data' } } }],
+        ['contentBlockDelta', { 'contentBlockIndex' => 0,
+                                'delta' => { 'toolUse' => { 'input' => '{"symbol":"MNQM26",' } } }],
+        ['contentBlockStart', { 'contentBlockIndex' => 1,
+                                'start' => { 'toolUse' => { 'toolUseId' => 'call_2', 'name' => 'search' } } }],
+        ['contentBlockDelta', { 'contentBlockIndex' => 1,
+                                'delta' => { 'toolUse' => { 'input' => '{"query":"market news"}' } } }],
+        # A delta for block 0 arriving after block 1's contentBlockStart — the exact
+        # ordering that corrupted arguments before nest_event_under_type's re-nested
+        # payload was routed by contentBlockIndex instead of falling back to nil.
+        ['contentBlockDelta', { 'contentBlockIndex' => 0,
+                                'delta' => { 'toolUse' => { 'input' => '"interval":"minute"}' } } }],
+        ['contentBlockStop', { 'contentBlockIndex' => 0 }],
+        ['contentBlockStop', { 'contentBlockIndex' => 1 }],
+        ['messageStop', { 'stopReason' => 'tool_use' }]
+      )
+
+      expect(message.tool_calls['call_1'].arguments).to eq('symbol' => 'MNQM26', 'interval' => 'minute')
+      expect(message.tool_calls['call_2'].arguments).to eq('query' => 'market news')
+    end
+
     it 'nests an exception payload under its :exception-type header' do
       message = Aws::EventStream::Message.new(
         headers: {

spec/ruby_llm/stream_accumulator_spec.rb

@@ -43,6 +43,29 @@
       )
     end
 
+    it 'raises ToolCallArgumentsTruncatedError instead of a bare JSON::ParserError for truncated arguments' do
+      accumulator = described_class.new
+
+      accumulator.add(RubyLLM::Chunk.new(
+                        role: :assistant, content: nil,
+                        tool_calls: { 'call_1' => RubyLLM::ToolCall.new(id: 'call_1', name: 'market_data',
+                                                                        arguments: {}) }
+                      ))
+      accumulator.add(RubyLLM::Chunk.new(
+                        role: :assistant, content: nil,
+                        tool_calls: { 'call_1' => RubyLLM::ToolCall.new(id: nil, name: nil,
+                                                                        arguments: '{"symbol":"MNQM26"') }
+                      ))
+      accumulator.add(RubyLLM::Chunk.new(role: :assistant, content: nil, finish_reason: 'tool_use'))
+
+      expect { accumulator.to_message(nil) }.to raise_error(RubyLLM::ToolCallArgumentsTruncatedError) do |error|
+        expect(error.tool_call_id).to eq('call_1')
+        expect(error.tool_name).to eq('market_data')
+        expect(error.raw_arguments).to eq('{"symbol":"MNQM26"')
+        expect(error.finish_reason).to eq('tool_use')
+      end
+    end
+
     it 'deduplicates citations repeated across chunks' do
       accumulator = described_class.new
       citation = RubyLLM::Citation.new(url: 'https://example.com', title: 'Example')