diff --git a/lib/ruby_llm/error.rb b/lib/ruby_llm/error.rb index 6f7ed01ec..693ece484 100644 --- a/lib/ruby_llm/error.rb +++ b/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 diff --git a/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb b/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb index 62b72bd8b..f9dc9f90f 100644 --- a/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb +++ b/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] diff --git a/lib/ruby_llm/protocols/converse/streaming.rb b/lib/ruby_llm/protocols/converse/streaming.rb index 431b1d7e4..39571e972 100644 --- a/lib/ruby_llm/protocols/converse/streaming.rb +++ b/lib/ruby_llm/protocols/converse/streaming.rb @@ -384,13 +384,20 @@ 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'] || {} @@ -398,11 +405,15 @@ def extract_tool_call_start(event) } 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) diff --git a/lib/ruby_llm/stream_accumulator.rb b/lib/ruby_llm/stream_accumulator.rb index 60eff437d..80c6bcbbc 100644 --- a/lib/ruby_llm/stream_accumulator.rb +++ b/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? diff --git a/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb b/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb index 02f5be6c3..33220993e 100644 --- a/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb +++ b/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 # --------------------------------------------------------------------------- diff --git a/spec/ruby_llm/protocols/converse/streaming_spec.rb b/spec/ruby_llm/protocols/converse/streaming_spec.rb index 1fdda8d71..639b22093 100644 --- a/spec/ruby_llm/protocols/converse/streaming_spec.rb +++ b/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: { diff --git a/spec/ruby_llm/stream_accumulator_spec.rb b/spec/ruby_llm/stream_accumulator_spec.rb index 307bdf5cc..7505fdc2a 100644 --- a/spec/ruby_llm/stream_accumulator_spec.rb +++ b/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')