bedrock: add error wrapping in StatusError for retry classification

Bedrock was the only provider not wrapping SDK errors in
*modelerrors.StatusError. This meant ClassifyModelError had to fall back
to regex parsing for Bedrock errors.

Add wrapBedrockError() which extracts the HTTP status code and
Retry-After header from smithyhttp.ResponseError, and call it from both
the stream adapter Recv() and CreateChatCompletionStream().

Assisted-By: docker-agent

Author: dgageot

go.mod

@@ -19,6 +19,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/credentials v1.19.12
 	github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.50.2
 	github.com/aws/aws-sdk-go-v2/service/sts v1.41.9
+	github.com/aws/smithy-go v1.24.2
 	github.com/aymanbagabas/go-udiff v0.4.1
 	github.com/blevesearch/bleve/v2 v2.5.7
 	github.com/bmatcuk/doublestar/v4 v4.10.0
@@ -90,7 +91,6 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/signin v1.0.8 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sso v1.30.13 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.17 // indirect
-	github.com/aws/smithy-go v1.24.2 // indirect
 	github.com/aymerick/douceur v0.2.0 // indirect
 	github.com/bahlo/generic-list-go v0.2.0 // indirect
 	github.com/bits-and-blooms/bitset v1.24.4 // indirect

pkg/model/provider/bedrock/adapter.go

@@ -71,7 +71,7 @@ func (a *streamAdapter) Recv() (chat.MessageStreamResponse, error) {
 		// Check for errors
 		if err := a.stream.Err(); err != nil {
 			slog.Debug("Bedrock stream: error on channel close", "error", err)
-			return chat.MessageStreamResponse{}, err
+			return chat.MessageStreamResponse{}, wrapBedrockError(err)
 		}
 		// If we have a pending finish reason but never got metadata, emit it now
 		if a.pendingFinishReason != "" {

pkg/model/provider/bedrock/client.go

@@ -219,7 +219,7 @@ func (c *Client) CreateChatCompletionStream(
 	output, err := c.bedrockClient.ConverseStream(ctx, input)
 	if err != nil {
 		slog.Error("Bedrock ConverseStream failed", "error", err)
-		return nil, fmt.Errorf("bedrock converse stream failed: %w", err)
+		return nil, wrapBedrockError(fmt.Errorf("bedrock converse stream failed: %w", err))
 	}
 
 	trackUsage := c.ModelConfig.TrackUsage == nil || *c.ModelConfig.TrackUsage

pkg/model/provider/bedrock/wrap.go

@@ -0,0 +1,35 @@
+package bedrock
+
+import (
+	"errors"
+
+	smithyhttp "github.com/aws/smithy-go/transport/http"
+
+	"github.com/docker/docker-agent/pkg/modelerrors"
+)
+
+// wrapBedrockError wraps an AWS Bedrock SDK error in a *modelerrors.StatusError
+// to carry HTTP status code metadata for the retry loop.
+// The AWS SDK v2 exposes HTTP status via smithyhttp.ResponseError.
+// Non-AWS errors (e.g., io.EOF, network errors) pass through unchanged.
+func wrapBedrockError(err error) error {
+	if err == nil {
+		return nil
+	}
+
+	var respErr *smithyhttp.ResponseError
+	if !errors.As(err, &respErr) {
+		return err
+	}
+
+	var resp *smithyhttp.Response
+	if respErr.HTTPResponse() != nil {
+		resp = respErr.HTTPResponse()
+	}
+
+	statusCode := respErr.HTTPStatusCode()
+	if resp != nil {
+		return modelerrors.WrapHTTPError(statusCode, resp.Response, err)
+	}
+	return modelerrors.WrapHTTPError(statusCode, nil, err)
+}

pkg/model/provider/bedrock/wrap_test.go

@@ -0,0 +1,110 @@
+package bedrock
+
+import (
+	"errors"
+	"fmt"
+	"net/http"
+	"testing"
+	"time"
+
+	smithy "github.com/aws/smithy-go"
+	smithyhttp "github.com/aws/smithy-go/transport/http"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/docker/docker-agent/pkg/modelerrors"
+)
+
+func makeTestBedrockError(statusCode int, retryAfterValue string) error {
+	header := http.Header{}
+	if retryAfterValue != "" {
+		header.Set("Retry-After", retryAfterValue)
+	}
+
+	httpResp := &http.Response{
+		StatusCode: statusCode,
+		Header:     header,
+	}
+	resp := &smithyhttp.Response{Response: httpResp}
+
+	return &smithy.OperationError{
+		ServiceID:     "BedrockRuntime",
+		OperationName: "ConverseStream",
+		Err: &smithyhttp.ResponseError{
+			Response: resp,
+			Err: &smithy.GenericAPIError{
+				Code:    "ThrottlingException",
+				Message: "Rate exceeded",
+			},
+		},
+	}
+}
+
+func TestWrapBedrockError(t *testing.T) {
+	t.Parallel()
+
+	t.Run("nil returns nil", func(t *testing.T) {
+		t.Parallel()
+		assert.NoError(t, wrapBedrockError(nil))
+	})
+
+	t.Run("non-AWS error passes through unchanged", func(t *testing.T) {
+		t.Parallel()
+		orig := errors.New("some network error")
+		result := wrapBedrockError(orig)
+		assert.Equal(t, orig, result)
+		var se *modelerrors.StatusError
+		assert.NotErrorAs(t, result, &se)
+	})
+
+	t.Run("429 without Retry-After wraps with zero RetryAfter", func(t *testing.T) {
+		t.Parallel()
+		awsErr := makeTestBedrockError(429, "")
+		result := wrapBedrockError(awsErr)
+		var se *modelerrors.StatusError
+		require.ErrorAs(t, result, &se)
+		assert.Equal(t, 429, se.StatusCode)
+		assert.Equal(t, time.Duration(0), se.RetryAfter)
+		// Original error still accessible
+		assert.ErrorIs(t, result, awsErr)
+	})
+
+	t.Run("429 with Retry-After header sets RetryAfter", func(t *testing.T) {
+		t.Parallel()
+		awsErr := makeTestBedrockError(429, "20")
+		result := wrapBedrockError(awsErr)
+		var se *modelerrors.StatusError
+		require.ErrorAs(t, result, &se)
+		assert.Equal(t, 429, se.StatusCode)
+		assert.Equal(t, 20*time.Second, se.RetryAfter)
+	})
+
+	t.Run("500 wraps with correct status code", func(t *testing.T) {
+		t.Parallel()
+		awsErr := makeTestBedrockError(500, "")
+		result := wrapBedrockError(awsErr)
+		var se *modelerrors.StatusError
+		require.ErrorAs(t, result, &se)
+		assert.Equal(t, 500, se.StatusCode)
+		assert.Equal(t, time.Duration(0), se.RetryAfter)
+	})
+
+	t.Run("wrapped error is classified correctly by ClassifyModelError", func(t *testing.T) {
+		t.Parallel()
+		awsErr := makeTestBedrockError(429, "15")
+		result := wrapBedrockError(awsErr)
+		retryable, rateLimited, retryAfter := modelerrors.ClassifyModelError(result)
+		assert.False(t, retryable)
+		assert.True(t, rateLimited)
+		assert.Equal(t, 15*time.Second, retryAfter)
+	})
+
+	t.Run("wrapped in fmt.Errorf still classified correctly", func(t *testing.T) {
+		t.Parallel()
+		awsErr := makeTestBedrockError(500, "")
+		wrapped := fmt.Errorf("bedrock converse stream failed: %w", wrapBedrockError(awsErr))
+		retryable, rateLimited, _ := modelerrors.ClassifyModelError(wrapped)
+		assert.True(t, retryable)
+		assert.False(t, rateLimited)
+	})
+}