From ad5232ade8035012facaab0f3fae5bb2406d8599 Mon Sep 17 00:00:00 2001 From: Andy Lo-A-Foe Date: Thu, 23 Apr 2026 23:27:55 +0200 Subject: [PATCH 1/2] feat(bedrock): implement StreamingProvider for real-time token streaming Adds ConverseStream API support to the Bedrock provider, implementing the StreamingProvider interface. Tokens flow via onChunk callback for real-time delivery to streaming-capable channels. - Extract buildConverseParams to share request logic between Chat and ChatStream - Add converseStreamReader interface for testability - Preserve raw payload in Arguments on JSON parse failure - Ensure Function.Arguments is always valid JSON - Streaming timeout only applied when explicitly configured - Capture stream Close() errors for diagnostics - Consistent "bedrock conversestream" / "bedrock:" log prefixes Co-Authored-By: Claude Opus 4.6 --- pkg/providers/bedrock/provider_bedrock.go | 284 ++++++++++++++++++---- 1 file changed, 239 insertions(+), 45 deletions(-) diff --git a/pkg/providers/bedrock/provider_bedrock.go b/pkg/providers/bedrock/provider_bedrock.go index 3798c5fd8..ee0ac75a0 100644 --- a/pkg/providers/bedrock/provider_bedrock.go +++ b/pkg/providers/bedrock/provider_bedrock.go @@ -135,48 +135,23 @@ func NewProvider(ctx context.Context, opts ...Option) (*Provider, error) { }, nil } -// Chat sends messages to AWS Bedrock using the Converse API. -func (p *Provider) Chat( - ctx context.Context, - messages []Message, - tools []ToolDefinition, - model string, - options map[string]any, -) (*LLMResponse, error) { - // Apply request timeout if context doesn't already have a deadline. - // Use explicit timeout if set, otherwise fall back to common default. - effectiveTimeout := p.requestTimeout - if effectiveTimeout <= 0 { - effectiveTimeout = common.DefaultRequestTimeout - } - if _, hasDeadline := ctx.Deadline(); !hasDeadline { - var cancel context.CancelFunc - ctx, cancel = context.WithTimeout(ctx, effectiveTimeout) - defer cancel() - } +// converseParams holds the shared request parameters for Converse and ConverseStream. +type converseParams struct { + messages []types.Message + system []types.SystemContentBlock + inferenceConfig *types.InferenceConfiguration + toolConfig *types.ToolConfiguration +} - // Build the Converse API input - input := &bedrockruntime.ConverseInput{ - ModelId: aws.String(model), - } - - // Convert messages to Bedrock format +func buildConverseParams(messages []Message, tools []ToolDefinition, options map[string]any) converseParams { bedrockMessages, systemPrompts := convertMessages(messages) - input.Messages = bedrockMessages - // Set system prompts if any - if len(systemPrompts) > 0 { - input.System = systemPrompts - } - - // Set inference configuration only when options are provided var inferenceConfig *types.InferenceConfiguration if maxTokens, ok := common.AsInt(options["max_tokens"]); ok && maxTokens > 0 { if inferenceConfig == nil { inferenceConfig = &types.InferenceConfiguration{} } - // Clamp to int32 range to avoid overflow if maxTokens > math.MaxInt32 { maxTokens = math.MaxInt32 } @@ -190,23 +165,53 @@ func (p *Provider) Chat( inferenceConfig.Temperature = aws.Float32(float32(temp)) } - if inferenceConfig != nil { - input.InferenceConfig = inferenceConfig - } - - // Convert tools to Bedrock format - // Only set ToolConfig if at least one valid tool was produced + var toolConfig *types.ToolConfiguration if len(tools) > 0 { - toolConfig := convertTools(tools) - if len(toolConfig.Tools) > 0 { - input.ToolConfig = toolConfig + tc := convertTools(tools) + if len(tc.Tools) > 0 { + toolConfig = tc } } - // Call Bedrock Converse API + return converseParams{ + messages: bedrockMessages, + system: systemPrompts, + inferenceConfig: inferenceConfig, + toolConfig: toolConfig, + } +} + +// Chat sends messages to AWS Bedrock using the Converse API. +func (p *Provider) Chat( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, +) (*LLMResponse, error) { + effectiveTimeout := p.requestTimeout + if effectiveTimeout <= 0 { + effectiveTimeout = common.DefaultRequestTimeout + } + if _, hasDeadline := ctx.Deadline(); !hasDeadline { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, effectiveTimeout) + defer cancel() + } + + params := buildConverseParams(messages, tools, options) + input := &bedrockruntime.ConverseInput{ + ModelId: aws.String(model), + Messages: params.messages, + InferenceConfig: params.inferenceConfig, + ToolConfig: params.toolConfig, + } + if len(params.system) > 0 { + input.System = params.system + } + output, err := p.client.Converse(ctx, input) if err != nil { - // Check for SSO token expiration errors and provide actionable guidance if isSSOTokenError(err) { return nil, fmt.Errorf( "bedrock converse: AWS credentials may have expired. If using AWS SSO, run 'aws sso login' to refresh: %w", @@ -216,10 +221,199 @@ func (p *Provider) Chat( return nil, fmt.Errorf("bedrock converse: %w", err) } - // Parse the response return parseResponse(output) } +// ChatStream sends messages to AWS Bedrock using the ConverseStream API. +// It streams the accumulated text so far via the onChunk callback and returns the complete response. +func (p *Provider) ChatStream( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, + onChunk func(accumulated string), +) (*LLMResponse, error) { + if p.requestTimeout > 0 { + if _, hasDeadline := ctx.Deadline(); !hasDeadline { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, p.requestTimeout) + defer cancel() + } + } + + params := buildConverseParams(messages, tools, options) + input := &bedrockruntime.ConverseStreamInput{ + ModelId: aws.String(model), + Messages: params.messages, + InferenceConfig: params.inferenceConfig, + ToolConfig: params.toolConfig, + } + if len(params.system) > 0 { + input.System = params.system + } + + output, err := p.client.ConverseStream(ctx, input) + if err != nil { + if isSSOTokenError(err) { + return nil, fmt.Errorf( + "bedrock conversestream: AWS credentials may have expired. If using AWS SSO, run 'aws sso login' to refresh: %w", + err, + ) + } + return nil, fmt.Errorf("bedrock conversestream: %w", err) + } + + return parseStreamResponse(ctx, output.GetStream(), onChunk) +} + +// converseStreamReader abstracts the Bedrock event stream so parseStreamResponse +// can be unit-tested with a mock event source. +type converseStreamReader interface { + Events() <-chan types.ConverseStreamOutput + Err() error + Close() error +} + +// parseStreamResponse processes the ConverseStream event stream and accumulates the response. +func parseStreamResponse( + ctx context.Context, + stream converseStreamReader, + onChunk func(accumulated string), +) (resp *LLMResponse, err error) { + if stream == nil { + return nil, fmt.Errorf("bedrock conversestream: nil event stream") + } + defer func() { + if closeErr := stream.Close(); closeErr != nil { + if err == nil { + err = fmt.Errorf("bedrock conversestream: close event stream: %w", closeErr) + } else { + log.Printf("bedrock conversestream: close event stream: %v", closeErr) + } + } + }() + + var textContent strings.Builder + finishReason := "stop" + var usage *UsageInfo + toolCalls := make([]ToolCall, 0) + + // Track active tool use blocks by index + type toolAccum struct { + id string + name string + argsJSON strings.Builder + } + activeTools := map[int]*toolAccum{} + + events := stream.Events() + for { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case event, ok := <-events: + if !ok { + // Stream closed + goto done + } + + switch e := event.(type) { + case *types.ConverseStreamOutputMemberContentBlockStart: + // New content block starting + if toolUse, ok := e.Value.Start.(*types.ContentBlockStartMemberToolUse); ok { + activeTools[int(aws.ToInt32(e.Value.ContentBlockIndex))] = &toolAccum{ + id: aws.ToString(toolUse.Value.ToolUseId), + name: aws.ToString(toolUse.Value.Name), + } + } + + case *types.ConverseStreamOutputMemberContentBlockDelta: + // Content delta + switch delta := e.Value.Delta.(type) { + case *types.ContentBlockDeltaMemberText: + textContent.WriteString(delta.Value) + if onChunk != nil { + onChunk(textContent.String()) + } + case *types.ContentBlockDeltaMemberToolUse: + idx := int(aws.ToInt32(e.Value.ContentBlockIndex)) + if tool, exists := activeTools[idx]; exists { + tool.argsJSON.WriteString(aws.ToString(delta.Value.Input)) + } + } + + case *types.ConverseStreamOutputMemberContentBlockStop: + // Content block finished - finalize tool if it was a tool use + idx := int(aws.ToInt32(e.Value.ContentBlockIndex)) + if tool, exists := activeTools[idx]; exists { + args := make(map[string]any) + argsStr := tool.argsJSON.String() + if argsStr != "" { + if err := json.Unmarshal([]byte(argsStr), &args); err != nil { + log.Printf("bedrock: stream: failed to parse tool arguments for %q: %v", tool.name, err) + args = map[string]any{"raw": argsStr} + } + } + funcArgs := argsStr + if argsJSON, marshalErr := json.Marshal(args); marshalErr == nil { + funcArgs = string(argsJSON) + } + toolCalls = append(toolCalls, ToolCall{ + ID: tool.id, + Name: tool.name, + Arguments: args, + Function: &FunctionCall{ + Name: tool.name, + Arguments: funcArgs, + }, + }) + delete(activeTools, idx) + } + + case *types.ConverseStreamOutputMemberMessageStop: + // Message complete + switch e.Value.StopReason { + case types.StopReasonToolUse: + finishReason = "tool_calls" + case types.StopReasonMaxTokens: + finishReason = "length" + case types.StopReasonEndTurn: + finishReason = "stop" + case types.StopReasonStopSequence: + finishReason = "stop" + case types.StopReasonContentFiltered: + finishReason = "content_filter" + default: + finishReason = "stop" + } + + case *types.ConverseStreamOutputMemberMetadata: + // Usage metadata + if e.Value.Usage != nil { + usage = &UsageInfo{ + PromptTokens: int(aws.ToInt32(e.Value.Usage.InputTokens)), + CompletionTokens: int(aws.ToInt32(e.Value.Usage.OutputTokens)), + TotalTokens: int(aws.ToInt32(e.Value.Usage.InputTokens)) + int(aws.ToInt32(e.Value.Usage.OutputTokens)), + } + } + } + } + } + +done: + if err := stream.Err(); err != nil { + return nil, fmt.Errorf("bedrock conversestream: %w", err) + } + + return &LLMResponse{ + Content: textContent.String(), + ToolCalls: toolCalls, + FinishReason: finishReason, + Usage: usage, + }, nil +} + // GetDefaultModel returns an empty string as Bedrock models are user-configured. func (p *Provider) GetDefaultModel() string { return "" From b03fa6176429ec5c81fca93c7818aa9dea6f308b Mon Sep 17 00:00:00 2001 From: Andy Lo-A-Foe Date: Thu, 23 Apr 2026 23:27:56 +0200 Subject: [PATCH 2/2] test(bedrock): add unit tests for ChatStream/parseStreamResponse Tests cover: text-only streaming with chunk accumulation, tool call parsing with fragmented JSON, mixed text+tool responses, context cancellation, invalid JSON fallback to raw payload, nil stream guard, default finish reason, and all stop reason mappings. Co-Authored-By: Claude Opus 4.6 --- .../bedrock/provider_bedrock_test.go | 270 ++++++++++++++++++ 1 file changed, 270 insertions(+) diff --git a/pkg/providers/bedrock/provider_bedrock_test.go b/pkg/providers/bedrock/provider_bedrock_test.go index 38a5e26da..9d6c747f1 100644 --- a/pkg/providers/bedrock/provider_bedrock_test.go +++ b/pkg/providers/bedrock/provider_bedrock_test.go @@ -8,6 +8,7 @@ package bedrock import ( + "context" "fmt" "testing" @@ -605,3 +606,272 @@ func TestIsSSOTokenError(t *testing.T) { }) } } + +// mockStreamReader implements bedrockruntime.ConverseStreamOutputReader for testing. +type mockStreamReader struct { + ch chan types.ConverseStreamOutput + err error +} + +func (r *mockStreamReader) Events() <-chan types.ConverseStreamOutput { return r.ch } +func (r *mockStreamReader) Close() error { return nil } +func (r *mockStreamReader) Err() error { return r.err } + +func newMockStream(events []types.ConverseStreamOutput) *bedrockruntime.ConverseStreamEventStream { + ch := make(chan types.ConverseStreamOutput, len(events)) + for _, e := range events { + ch <- e + } + close(ch) + + return bedrockruntime.NewConverseStreamEventStream(func(es *bedrockruntime.ConverseStreamEventStream) { + es.Reader = &mockStreamReader{ch: ch} + }) +} + +func TestParseStreamResponse_TextOnly(t *testing.T) { + events := []types.ConverseStreamOutput{ + &types.ConverseStreamOutputMemberContentBlockDelta{ + Value: types.ContentBlockDeltaEvent{ + Delta: &types.ContentBlockDeltaMemberText{Value: "Hello "}, + ContentBlockIndex: aws.Int32(0), + }, + }, + &types.ConverseStreamOutputMemberContentBlockDelta{ + Value: types.ContentBlockDeltaEvent{ + Delta: &types.ContentBlockDeltaMemberText{Value: "World"}, + ContentBlockIndex: aws.Int32(0), + }, + }, + &types.ConverseStreamOutputMemberMessageStop{ + Value: types.MessageStopEvent{StopReason: types.StopReasonEndTurn}, + }, + &types.ConverseStreamOutputMemberMetadata{ + Value: types.ConverseStreamMetadataEvent{ + Usage: &types.TokenUsage{ + InputTokens: aws.Int32(10), + OutputTokens: aws.Int32(5), + }, + }, + }, + } + + var chunks []string + stream := newMockStream(events) + resp, err := parseStreamResponse(context.Background(), stream, func(accumulated string) { + chunks = append(chunks, accumulated) + }) + + require.NoError(t, err) + assert.Equal(t, "Hello World", resp.Content) + assert.Equal(t, "stop", resp.FinishReason) + assert.Empty(t, resp.ToolCalls) + require.NotNil(t, resp.Usage) + assert.Equal(t, 10, resp.Usage.PromptTokens) + assert.Equal(t, 5, resp.Usage.CompletionTokens) + assert.Equal(t, 15, resp.Usage.TotalTokens) + assert.Equal(t, []string{"Hello ", "Hello World"}, chunks) +} + +func TestParseStreamResponse_ToolCall(t *testing.T) { + events := []types.ConverseStreamOutput{ + &types.ConverseStreamOutputMemberContentBlockStart{ + Value: types.ContentBlockStartEvent{ + ContentBlockIndex: aws.Int32(0), + Start: &types.ContentBlockStartMemberToolUse{ + Value: types.ToolUseBlockStart{ + ToolUseId: aws.String("call_1"), + Name: aws.String("search"), + }, + }, + }, + }, + &types.ConverseStreamOutputMemberContentBlockDelta{ + Value: types.ContentBlockDeltaEvent{ + ContentBlockIndex: aws.Int32(0), + Delta: &types.ContentBlockDeltaMemberToolUse{ + Value: types.ToolUseBlockDelta{Input: aws.String(`{"q":`)}, + }, + }, + }, + &types.ConverseStreamOutputMemberContentBlockDelta{ + Value: types.ContentBlockDeltaEvent{ + ContentBlockIndex: aws.Int32(0), + Delta: &types.ContentBlockDeltaMemberToolUse{ + Value: types.ToolUseBlockDelta{Input: aws.String(`"test"}`)}, + }, + }, + }, + &types.ConverseStreamOutputMemberContentBlockStop{ + Value: types.ContentBlockStopEvent{ContentBlockIndex: aws.Int32(0)}, + }, + &types.ConverseStreamOutputMemberMessageStop{ + Value: types.MessageStopEvent{StopReason: types.StopReasonToolUse}, + }, + } + + stream := newMockStream(events) + resp, err := parseStreamResponse(context.Background(), stream, nil) + + require.NoError(t, err) + assert.Equal(t, "tool_calls", resp.FinishReason) + require.Len(t, resp.ToolCalls, 1) + assert.Equal(t, "call_1", resp.ToolCalls[0].ID) + assert.Equal(t, "search", resp.ToolCalls[0].Name) + assert.Equal(t, map[string]any{"q": "test"}, resp.ToolCalls[0].Arguments) + require.NotNil(t, resp.ToolCalls[0].Function) + assert.Equal(t, "search", resp.ToolCalls[0].Function.Name) + assert.Equal(t, `{"q":"test"}`, resp.ToolCalls[0].Function.Arguments) +} + +func TestParseStreamResponse_TextAndToolCall(t *testing.T) { + events := []types.ConverseStreamOutput{ + &types.ConverseStreamOutputMemberContentBlockDelta{ + Value: types.ContentBlockDeltaEvent{ + ContentBlockIndex: aws.Int32(0), + Delta: &types.ContentBlockDeltaMemberText{Value: "Let me search that."}, + }, + }, + &types.ConverseStreamOutputMemberContentBlockStart{ + Value: types.ContentBlockStartEvent{ + ContentBlockIndex: aws.Int32(1), + Start: &types.ContentBlockStartMemberToolUse{ + Value: types.ToolUseBlockStart{ + ToolUseId: aws.String("call_2"), + Name: aws.String("web"), + }, + }, + }, + }, + &types.ConverseStreamOutputMemberContentBlockDelta{ + Value: types.ContentBlockDeltaEvent{ + ContentBlockIndex: aws.Int32(1), + Delta: &types.ContentBlockDeltaMemberToolUse{ + Value: types.ToolUseBlockDelta{Input: aws.String(`{"url":"https://example.com"}`)}, + }, + }, + }, + &types.ConverseStreamOutputMemberContentBlockStop{ + Value: types.ContentBlockStopEvent{ContentBlockIndex: aws.Int32(1)}, + }, + &types.ConverseStreamOutputMemberMessageStop{ + Value: types.MessageStopEvent{StopReason: types.StopReasonToolUse}, + }, + } + + var chunks []string + stream := newMockStream(events) + resp, err := parseStreamResponse(context.Background(), stream, func(accumulated string) { + chunks = append(chunks, accumulated) + }) + + require.NoError(t, err) + assert.Equal(t, "Let me search that.", resp.Content) + assert.Equal(t, "tool_calls", resp.FinishReason) + require.Len(t, resp.ToolCalls, 1) + assert.Equal(t, "web", resp.ToolCalls[0].Name) + assert.Equal(t, []string{"Let me search that."}, chunks) +} + +func TestParseStreamResponse_ContextCancelled(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + // Use an unbuffered channel with no events so ctx.Done() is the only ready case. + ch := make(chan types.ConverseStreamOutput) + + stream := bedrockruntime.NewConverseStreamEventStream(func(es *bedrockruntime.ConverseStreamEventStream) { + es.Reader = &mockStreamReader{ch: ch} + }) + + _, err := parseStreamResponse(ctx, stream, nil) + assert.ErrorIs(t, err, context.Canceled) +} + +func TestParseStreamResponse_InvalidToolJSON(t *testing.T) { + events := []types.ConverseStreamOutput{ + &types.ConverseStreamOutputMemberContentBlockStart{ + Value: types.ContentBlockStartEvent{ + ContentBlockIndex: aws.Int32(0), + Start: &types.ContentBlockStartMemberToolUse{ + Value: types.ToolUseBlockStart{ + ToolUseId: aws.String("call_bad"), + Name: aws.String("broken"), + }, + }, + }, + }, + &types.ConverseStreamOutputMemberContentBlockDelta{ + Value: types.ContentBlockDeltaEvent{ + ContentBlockIndex: aws.Int32(0), + Delta: &types.ContentBlockDeltaMemberToolUse{ + Value: types.ToolUseBlockDelta{Input: aws.String(`{not valid json`)}, + }, + }, + }, + &types.ConverseStreamOutputMemberContentBlockStop{ + Value: types.ContentBlockStopEvent{ContentBlockIndex: aws.Int32(0)}, + }, + &types.ConverseStreamOutputMemberMessageStop{ + Value: types.MessageStopEvent{StopReason: types.StopReasonToolUse}, + }, + } + + stream := newMockStream(events) + resp, err := parseStreamResponse(context.Background(), stream, nil) + + require.NoError(t, err) + require.Len(t, resp.ToolCalls, 1) + assert.Equal(t, map[string]any{"raw": `{not valid json`}, resp.ToolCalls[0].Arguments) + assert.JSONEq(t, `{"raw":"{not valid json"}`, resp.ToolCalls[0].Function.Arguments) +} + +func TestParseStreamResponse_DefaultFinishReason(t *testing.T) { + events := []types.ConverseStreamOutput{ + &types.ConverseStreamOutputMemberContentBlockDelta{ + Value: types.ContentBlockDeltaEvent{ + Delta: &types.ContentBlockDeltaMemberText{Value: "partial"}, + ContentBlockIndex: aws.Int32(0), + }, + }, + } + + stream := newMockStream(events) + resp, err := parseStreamResponse(context.Background(), stream, nil) + + require.NoError(t, err) + assert.Equal(t, "stop", resp.FinishReason) +} + +func TestParseStreamResponse_NilStream(t *testing.T) { + _, err := parseStreamResponse(context.Background(), nil, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "nil event stream") +} + +func TestParseStreamResponse_StopReasons(t *testing.T) { + tests := []struct { + reason types.StopReason + expected string + }{ + {types.StopReasonEndTurn, "stop"}, + {types.StopReasonMaxTokens, "length"}, + {types.StopReasonToolUse, "tool_calls"}, + {types.StopReasonStopSequence, "stop"}, + {types.StopReasonContentFiltered, "content_filter"}, + } + + for _, tt := range tests { + t.Run(string(tt.reason), func(t *testing.T) { + events := []types.ConverseStreamOutput{ + &types.ConverseStreamOutputMemberMessageStop{ + Value: types.MessageStopEvent{StopReason: tt.reason}, + }, + } + stream := newMockStream(events) + resp, err := parseStreamResponse(context.Background(), stream, nil) + require.NoError(t, err) + assert.Equal(t, tt.expected, resp.FinishReason) + }) + } +}