diff --git a/components/model/claude/bedrock_http_client_test.go b/components/model/claude/bedrock_http_client_test.go new file mode 100644 index 000000000..f59d4e7bd --- /dev/null +++ b/components/model/claude/bedrock_http_client_test.go @@ -0,0 +1,341 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package claude + +import ( + "context" + "crypto/tls" + "encoding/pem" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +func TestNewBedrockAWSHTTPClient(t *testing.T) { + t.Run("nil client", func(t *testing.T) { + assert.Nil(t, newBedrockAWSHTTPClient(nil)) + }) + + t.Run("custom RoundTripper", func(t *testing.T) { + client := &http.Client{ + Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, nil + }), + Timeout: 2 * time.Second, + } + awsClient := newBedrockAWSHTTPClient(client) + if !assert.NotNil(t, awsClient) { + return + } + assert.Equal(t, client.Timeout, awsClient.GetTimeout()) + }) + + t.Run("standard transport", func(t *testing.T) { + transport := http.DefaultTransport.(*http.Transport).Clone() + transport.DisableCompression = true + transport.MaxIdleConns = 37 + transport.TLSClientConfig = &tls.Config{MinVersion: tls.VersionTLS13} + client := &http.Client{ + Transport: transport, + Timeout: 3 * time.Second, + } + + awsClient := newBedrockAWSHTTPClient(client) + if !assert.NotNil(t, awsClient) { + return + } + assert.Equal(t, client.Timeout, awsClient.GetTimeout()) + + gotTransport := awsClient.GetTransport() + assert.NotSame(t, transport, gotTransport) + assert.Equal(t, transport.DisableCompression, gotTransport.DisableCompression) + assert.Equal(t, transport.MaxIdleConns, gotTransport.MaxIdleConns) + assert.NotSame(t, transport.TLSClientConfig, gotTransport.TLSClientConfig) + assert.Equal(t, transport.TLSClientConfig.MinVersion, gotTransport.TLSClientConfig.MinVersion) + }) +} + +func TestNewChatModelBedrockHTTPClientWithCABundle(t *testing.T) { + setBedrockTestEnv(t) + t.Setenv("AWS_CA_BUNDLE", writeTestCABundle(t)) + + client := &http.Client{ + Transport: http.DefaultTransport.(*http.Transport).Clone(), + Timeout: 3 * time.Second, + } + + var ( + chatModel *ChatModel + err error + ) + assert.NotPanics(t, func() { + chatModel, err = NewChatModel(context.Background(), newBedrockTestConfig(client)) + }) + assert.NoError(t, err) + assert.NotNil(t, chatModel) +} + +func TestNewChatModelBedrockInvalidCABundleReturnsError(t *testing.T) { + setBedrockTestEnv(t) + bundlePath := filepath.Join(t.TempDir(), "invalid-ca-bundle.pem") + if err := os.WriteFile(bundlePath, []byte("not a certificate"), 0o600); err != nil { + t.Fatalf("write invalid CA bundle: %v", err) + } + t.Setenv("AWS_CA_BUNDLE", bundlePath) + + var err error + assert.NotPanics(t, func() { + _, err = NewChatModel(context.Background(), newBedrockTestConfig(&http.Client{})) + }) + assert.ErrorContains(t, err, "load AWS config for Bedrock") + assert.ErrorContains(t, err, "failed to load custom CA bundle PEM file") +} + +func TestBedrockRequestUsesCustomHTTPClientWithCABundle(t *testing.T) { + setBedrockTestEnv(t) + t.Setenv("AWS_CA_BUNDLE", writeTestCABundle(t)) + + originalDefaultTransport := http.DefaultTransport + fallbackRequestCount := 0 + http.DefaultTransport = roundTripFunc(func(req *http.Request) (*http.Response, error) { + fallbackRequestCount++ + return &http.Response{ + StatusCode: http.StatusTeapot, + Status: "418 I'm a teapot", + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader("unexpected default transport request")), + Request: req, + }, nil + }) + t.Cleanup(func() { + http.DefaultTransport = originalDefaultTransport + }) + + var ( + requestCount int + requestHost string + requestPath string + authorization string + ) + client := &http.Client{ + Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + requestCount++ + requestHost = req.URL.Host + requestPath = req.URL.Path + authorization = req.Header.Get("Authorization") + return &http.Response{ + StatusCode: http.StatusOK, + Status: "200 OK", + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{ + "id":"msg_test", + "type":"message", + "role":"assistant", + "content":[{"type":"text","text":"hi"}], + "model":"test-model", + "stop_reason":"end_turn", + "stop_sequence":null, + "usage":{"input_tokens":1,"output_tokens":1} + }`)), + Request: req, + }, nil + }), + Timeout: 3 * time.Second, + } + + chatModel, err := NewChatModel(context.Background(), newBedrockTestConfig(client)) + if !assert.NoError(t, err) { + return + } + + message, err := chatModel.Generate(context.Background(), []*schema.Message{schema.UserMessage("hello")}) + if !assert.NoError(t, err) { + return + } + assert.Equal(t, "hi", message.Content) + assert.Equal(t, 1, requestCount) + assert.Equal(t, "bedrock-runtime.us-east-1.amazonaws.com", requestHost) + assert.Equal(t, "/model/test-model/invoke", requestPath) + assert.True(t, strings.HasPrefix(authorization, "AWS4-HMAC-SHA256")) + assert.Zero(t, fallbackRequestCount) +} + +func TestBedrockCustomHTTPClientIsUsedForCredentialsWithoutCABundle(t *testing.T) { + setBedrockTestEnv(t) + setBedrockWebIdentityEnv(t) + t.Setenv("AWS_CA_BUNDLE", "") + + var ( + stsRequestCount int + bedrockRequestCount int + bedrockAuthorization string + ) + client := &http.Client{ + Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + switch { + case strings.HasPrefix(req.URL.Host, "sts."): + stsRequestCount++ + return &http.Response{ + StatusCode: http.StatusOK, + Status: "200 OK", + Header: http.Header{"Content-Type": []string{"text/xml"}}, + Body: io.NopCloser(strings.NewReader(` + + + arn:aws:sts::123456789012:assumed-role/test-role/test-session + test-role-id:test-session + + + web-identity-access-key + web-identity-secret-key + web-identity-session-token + 2100-01-01T00:00:00Z + + + request-id + `)), + Request: req, + }, nil + case req.URL.Path == "/model/test-model/invoke": + bedrockRequestCount++ + bedrockAuthorization = req.Header.Get("Authorization") + return &http.Response{ + StatusCode: http.StatusOK, + Status: "200 OK", + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{ + "id":"msg_test", + "type":"message", + "role":"assistant", + "content":[{"type":"text","text":"credential path ok"}], + "model":"test-model", + "stop_reason":"end_turn", + "stop_sequence":null, + "usage":{"input_tokens":1,"output_tokens":1} + }`)), + Request: req, + }, nil + default: + t.Fatalf("unexpected request through custom HTTP client: %s", req.URL) + return nil, nil + } + }), + Timeout: 3 * time.Second, + } + + config := newBedrockTestConfig(client) + config.AccessKey = "" + config.SecretAccessKey = "" + chatModel, err := NewChatModel(context.Background(), config) + if !assert.NoError(t, err) { + return + } + + message, err := chatModel.Generate(context.Background(), []*schema.Message{schema.UserMessage("hello")}) + if !assert.NoError(t, err) { + return + } + assert.Equal(t, "credential path ok", message.Content) + assert.Equal(t, 1, stsRequestCount) + assert.Equal(t, 1, bedrockRequestCount) + assert.Contains(t, bedrockAuthorization, "Credential=web-identity-access-key/") +} + +func newBedrockTestConfig(client *http.Client) *Config { + return &Config{ + ByBedrock: true, + AccessKey: "test-access-key", + SecretAccessKey: "test-secret-key", + Region: "us-east-1", + Model: "test-model", + MaxTokens: 8, + HTTPClient: client, + } +} + +func setBedrockTestEnv(t *testing.T) { + t.Helper() + t.Setenv("ANTHROPIC_API_KEY", "") + t.Setenv("AWS_BEARER_TOKEN_BEDROCK", "") + t.Setenv("AWS_EC2_METADATA_DISABLED", "true") +} + +func setBedrockWebIdentityEnv(t *testing.T) { + t.Helper() + + for _, key := range []string{ + "AWS_ACCESS_KEY_ID", + "AWS_ACCESS_KEY", + "AWS_SECRET_ACCESS_KEY", + "AWS_SECRET_KEY", + "AWS_SESSION_TOKEN", + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", + "AWS_CONTAINER_CREDENTIALS_FULL_URI", + } { + t.Setenv(key, "") + } + + configPath := filepath.Join(t.TempDir(), "config") + credentialsPath := filepath.Join(t.TempDir(), "credentials") + for _, path := range []string{configPath, credentialsPath} { + if err := os.WriteFile(path, []byte("[default]\n"), 0o600); err != nil { + t.Fatalf("write empty AWS config: %v", err) + } + } + t.Setenv("AWS_CONFIG_FILE", configPath) + t.Setenv("AWS_SHARED_CREDENTIALS_FILE", credentialsPath) + + tokenPath := filepath.Join(t.TempDir(), "web-identity-token") + if err := os.WriteFile(tokenPath, []byte("test-token"), 0o600); err != nil { + t.Fatalf("write web identity token: %v", err) + } + t.Setenv("AWS_WEB_IDENTITY_TOKEN_FILE", tokenPath) + t.Setenv("AWS_ROLE_ARN", "arn:aws:iam::123456789012:role/test-role") + t.Setenv("AWS_ROLE_SESSION_NAME", "test-session") +} + +func writeTestCABundle(t *testing.T) string { + t.Helper() + + server := httptest.NewTLSServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + certificate := server.Certificate() + server.Close() + + bundle := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certificate.Raw}) + path := filepath.Join(t.TempDir(), "ca-bundle.pem") + if err := os.WriteFile(path, bundle, 0o600); err != nil { + t.Fatalf("write CA bundle: %v", err) + } + return path +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} diff --git a/components/model/claude/claude.go b/components/model/claude/claude.go index c5e607095..07454c809 100644 --- a/components/model/claude/claude.go +++ b/components/model/claude/claude.go @@ -33,10 +33,12 @@ import ( "github.com/anthropics/anthropic-sdk-go/option" "github.com/anthropics/anthropic-sdk-go/packages/param" "github.com/anthropics/anthropic-sdk-go/vertex" + "github.com/aws/aws-sdk-go-v2/aws" + awshttp "github.com/aws/aws-sdk-go-v2/aws/transport/http" awsConfig "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/credentials" - "golang.org/x/oauth2/google" "github.com/eino-contrib/jsonschema" + "golang.org/x/oauth2/google" "github.com/cloudwego/eino/components" @@ -112,10 +114,16 @@ func NewChatModel(ctx context.Context, config *Config) (*ChatModel, error) { opts = append(opts, awsConfig.WithSharedConfigProfile(config.Profile)) } + awsCfg, err := loadBedrockAWSConfig(ctx, opts, config.HTTPClient) + if err != nil { + return nil, fmt.Errorf("load AWS config for Bedrock: %w", err) + } + + clientOpts := []option.RequestOption{bedrock.WithConfig(awsCfg)} if config.HTTPClient != nil { - opts = append(opts, awsConfig.WithHTTPClient(config.HTTPClient)) + clientOpts = append(clientOpts, option.WithHTTPClient(config.HTTPClient)) } - cli = anthropic.NewClient(bedrock.WithLoadDefaultConfig(ctx, opts...)) + cli = anthropic.NewClient(clientOpts...) } else { // Use direct Anthropic API var opts []option.RequestOption @@ -175,6 +183,59 @@ func NewChatModel(ctx context.Context, config *Config) (*ChatModel, error) { }, nil } +const unsupportedCABundleHTTPClientError = "unable to add custom RootCAs HTTPClient" + +// loadBedrockAWSConfig preserves the caller's HTTP client for credential and +// config requests whenever the AWS loader can use it directly. The loader can +// only add a custom CA bundle to its own BuildableClient, so retry with a +// converted client for that specific incompatibility. +func loadBedrockAWSConfig(ctx context.Context, opts []func(*awsConfig.LoadOptions) error, + client *http.Client) (aws.Config, error) { + if client == nil { + return awsConfig.LoadDefaultConfig(ctx, opts...) + } + + awsCfg, err := awsConfig.LoadDefaultConfig(ctx, withBedrockAWSHTTPClient(opts, client)...) + if err == nil || !strings.Contains(err.Error(), unsupportedCABundleHTTPClientError) { + return awsCfg, err + } + + return awsConfig.LoadDefaultConfig(ctx, withBedrockAWSHTTPClient(opts, newBedrockAWSHTTPClient(client))...) +} + +func withBedrockAWSHTTPClient(opts []func(*awsConfig.LoadOptions) error, + client awsConfig.HTTPClient) []func(*awsConfig.LoadOptions) error { + result := make([]func(*awsConfig.LoadOptions) error, len(opts), len(opts)+1) + copy(result, opts) + return append(result, awsConfig.WithHTTPClient(client)) +} + +// newBedrockAWSHTTPClient converts the standard HTTP transport into the +// buildable client required by the AWS config loader when AWS_CA_BUNDLE is set. +// The original client is still passed to the Anthropic SDK for Bedrock API +// requests. A custom RoundTripper cannot be converted without changing its +// behavior, so the AWS config loader gets a buildable default that only +// preserves the caller's timeout in that case. +func newBedrockAWSHTTPClient(client *http.Client) *awshttp.BuildableClient { + if client == nil { + return nil + } + + awsHTTPClient := awshttp.NewBuildableClient().WithTimeout(client.Timeout) + transport := client.Transport + if transport == nil { + transport = http.DefaultTransport + } + httpTransport, ok := transport.(*http.Transport) + if !ok { + return awsHTTPClient + } + + return awsHTTPClient.WithTransportOptions(func(transport *http.Transport) { + *transport = *httpTransport.Clone() + }) +} + // Config contains the configuration options for the Claude model type Config struct { // ByBedrock indicates whether to use Bedrock Service diff --git a/components/model/claude/go.mod b/components/model/claude/go.mod index 55a2e002e..54abf8af8 100644 --- a/components/model/claude/go.mod +++ b/components/model/claude/go.mod @@ -4,6 +4,7 @@ go 1.24 require ( github.com/anthropics/anthropic-sdk-go v1.56.0 + github.com/aws/aws-sdk-go-v2 v1.33.0 github.com/aws/aws-sdk-go-v2/config v1.29.1 github.com/aws/aws-sdk-go-v2/credentials v1.17.54 github.com/bytedance/mockey v1.2.13 @@ -17,7 +18,6 @@ require ( cloud.google.com/go/auth v0.7.2 // indirect cloud.google.com/go/auth/oauth2adapt v0.2.3 // indirect cloud.google.com/go/compute/metadata v0.5.0 // indirect - github.com/aws/aws-sdk-go-v2 v1.33.0 // indirect github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.3 // indirect github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.24 // indirect github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.28 // indirect