From d8878a5bcc1da9a209c5acd724d4c2aa6c34fa23 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 14 Jul 2026 14:34:29 +0000 Subject: [PATCH 1/6] fix: validate bedrock protocol at provider construction --- aibridge/provider/anthropic.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index 9424d4c9601..cc91d8d0efb 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -69,6 +69,9 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. // so it is cheap to run at construction. var bedrock *messages.BedrockRuntime if bedrockCfg != nil { + if err := bedrockCfg.Validate(); err != nil { + return nil, xerrors.Errorf("bedrock config: %w", err) + } creds, resolvedRegion, err := buildBedrockCredentials(ctx, *bedrockCfg) if err != nil { return nil, xerrors.Errorf("build bedrock credentials: %w", err) From f55d7147dd5523f0f88d002b1e6615e2cbcbddc3 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 14 Jul 2026 16:44:54 +0000 Subject: [PATCH 2/6] feat: add bedrock protocol to trace attributes --- aibridge/config/config.go | 10 ++++++++++ aibridge/intercept/messages/base.go | 6 +++++- aibridge/tracing/tracing.go | 13 +++++++------ 3 files changed, 22 insertions(+), 7 deletions(-) diff --git a/aibridge/config/config.go b/aibridge/config/config.go index acc1192319f..a3ba08f470c 100644 --- a/aibridge/config/config.go +++ b/aibridge/config/config.go @@ -70,6 +70,16 @@ type AWSBedrock struct { Protocol BedrockProtocol } +// ResolvedProtocol returns the configured protocol, mapping the empty value to +// the legacy InvokeModel protocol so existing providers keep the legacy +// behavior. +func (c AWSBedrock) ResolvedProtocol() BedrockProtocol { + if c.Protocol == "" { + return BedrockProtocolInvokeModel + } + return c.Protocol +} + // Validate verifies protocol-specific Bedrock configuration. func (c AWSBedrock) Validate() error { switch c.Protocol { diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index a0025d6881b..20f581bb59b 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -164,7 +164,7 @@ func (i *interceptionBase) Model() string { } func (i *interceptionBase) baseTraceAttributes(r *http.Request, streaming bool) []attribute.KeyValue { - return []attribute.KeyValue{ + attrs := []attribute.KeyValue{ attribute.String(tracing.RequestPath, r.URL.Path), attribute.String(tracing.InterceptionID, i.id.String()), attribute.String(tracing.InitiatorID, aibcontext.ActorIDFromContext(r.Context())), @@ -173,6 +173,10 @@ func (i *interceptionBase) baseTraceAttributes(r *http.Request, streaming bool) attribute.Bool(tracing.Streaming, streaming), attribute.Bool(tracing.IsBedrock, i.bedrock != nil), } + if i.bedrock != nil { + attrs = append(attrs, attribute.String(tracing.BedrockProtocol, string(i.bedrock.Cfg.ResolvedProtocol()))) + } + return attrs } func (i *interceptionBase) injectTools() { diff --git a/aibridge/tracing/tracing.go b/aibridge/tracing/tracing.go index 7adaf3f65e3..f01e3f1f330 100644 --- a/aibridge/tracing/tracing.go +++ b/aibridge/tracing/tracing.go @@ -17,12 +17,13 @@ const ( // trace attribute key constants RequestPath = "request_path" - InterceptionID = "interception_id" - InitiatorID = "user_id" - Provider = "provider" - Model = "model" - Streaming = "streaming" - IsBedrock = "aws_bedrock" + InterceptionID = "interception_id" + InitiatorID = "user_id" + Provider = "provider" + Model = "model" + Streaming = "streaming" + IsBedrock = "aws_bedrock" + BedrockProtocol = "aws_bedrock_protocol" PassthroughURL = "passthrough_url" PassthroughUpstreamURL = "passthrough_upstream_url" From c31a5c1244d791b59f118f7b1f239b63951f9a2b Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 14 Jul 2026 17:36:42 +0000 Subject: [PATCH 3/6] test: fix tests --- aibridge/internal/integrationtest/trace_internal_test.go | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/aibridge/internal/integrationtest/trace_internal_test.go b/aibridge/internal/integrationtest/trace_internal_test.go index 23fcbdeae8b..43719c3f177 100644 --- a/aibridge/internal/integrationtest/trace_internal_test.go +++ b/aibridge/internal/integrationtest/trace_internal_test.go @@ -187,6 +187,9 @@ func TestTraceAnthropic(t *testing.T) { attribute.Bool(tracing.Streaming, tc.streaming), attribute.Bool(tracing.IsBedrock, tc.bedrock), } + if tc.bedrock { + attrs = append(attrs, attribute.String(tracing.BedrockProtocol, string(config.BedrockProtocolInvokeModel))) + } require.Len(t, sr.Ended(), totalCount) verifyTraces(t, sr, tc.expect, attrs) @@ -311,6 +314,9 @@ func TestTraceAnthropicErr(t *testing.T) { attribute.Bool(tracing.Streaming, tc.streaming), attribute.Bool(tracing.IsBedrock, tc.bedrock), } + if tc.bedrock { + attrs = append(attrs, attribute.String(tracing.BedrockProtocol, string(config.BedrockProtocolInvokeModel))) + } verifyTraces(t, sr, tc.expect, attrs) }) @@ -422,6 +428,9 @@ func TestInjectedToolsTrace(t *testing.T) { } if tc.expectProvider == config.ProviderAnthropic { attrs = append(attrs, attribute.Bool(tracing.IsBedrock, tc.bedrock)) + if tc.bedrock { + attrs = append(attrs, attribute.String(tracing.BedrockProtocol, string(config.BedrockProtocolInvokeModel))) + } } verifyTraces(t, sr, []expectTrace{{"Intercept.ProcessRequest.ToolCall", 1, codes.Unset}}, attrs) From 8be08b289c0224d3b4e42a37f8dc03a03911a001 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 14 Jul 2026 17:43:28 +0000 Subject: [PATCH 4/6] fix: CR's fixes --- aibridge/provider/anthropic.go | 6 ++-- aibridge/provider/anthropic_internal_test.go | 36 ++++++++++++++++++++ 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index cc91d8d0efb..460b38102cd 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -69,9 +69,6 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. // so it is cheap to run at construction. var bedrock *messages.BedrockRuntime if bedrockCfg != nil { - if err := bedrockCfg.Validate(); err != nil { - return nil, xerrors.Errorf("bedrock config: %w", err) - } creds, resolvedRegion, err := buildBedrockCredentials(ctx, *bedrockCfg) if err != nil { return nil, xerrors.Errorf("build bedrock credentials: %w", err) @@ -82,6 +79,9 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. if runtimeCfg.Region == "" { runtimeCfg.Region = resolvedRegion } + if err := runtimeCfg.Validate(); err != nil { + return nil, xerrors.Errorf("bedrock config: %w", err) + } bedrock = &messages.BedrockRuntime{Cfg: runtimeCfg, Creds: creds} } diff --git a/aibridge/provider/anthropic_internal_test.go b/aibridge/provider/anthropic_internal_test.go index 9df6f037184..cdc8afe9148 100644 --- a/aibridge/provider/anthropic_internal_test.go +++ b/aibridge/provider/anthropic_internal_test.go @@ -114,6 +114,42 @@ func TestNewAnthropic_KeyResolution(t *testing.T) { } } +// NOTE: no t.Parallel() because the subtests use t.Setenv. +func TestNewAnthropic_BedrockRegionResolution(t *testing.T) { + t.Run("mantle_region_from_env", func(t *testing.T) { + t.Setenv("AWS_REGION", "us-west-2") + + p, err := NewAnthropic(context.Background(), config.Anthropic{}, &config.AWSBedrock{ + BaseURL: "https://bedrock-mantle.us-west-2.api.aws/anthropic", + Protocol: config.BedrockProtocolMantle, + AccessKey: "test-key", + AccessKeySecret: "test-secret", + }) + require.NoError(t, err) + require.NotNil(t, p.bedrock) + require.Equal(t, "us-west-2", p.bedrock.Cfg.Region) + }) + + t.Run("mantle_no_region_anywhere", func(t *testing.T) { + // Clear every source the AWS SDK consults for a region so none + // resolves, then confirm construction rejects the mantle provider. + t.Setenv("AWS_REGION", "") + t.Setenv("AWS_DEFAULT_REGION", "") + t.Setenv("AWS_PROFILE", "") + t.Setenv("AWS_CONFIG_FILE", "/dev/null") + t.Setenv("AWS_SHARED_CREDENTIALS_FILE", "/dev/null") + t.Setenv("AWS_EC2_METADATA_DISABLED", "true") + + _, err := NewAnthropic(context.Background(), config.Anthropic{}, &config.AWSBedrock{ + BaseURL: "https://proxy.internal", + Protocol: config.BedrockProtocolMantle, + AccessKey: "test-key", + AccessKeySecret: "test-secret", + }) + require.ErrorContains(t, err, "region required") + }) +} + func TestAnthropic_CreateInterceptor(t *testing.T) { t.Parallel() From 41add5c150812c4d34e32352d4b2a044c299e3ad Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 14 Jul 2026 18:10:29 +0000 Subject: [PATCH 5/6] test: fix tests --- cli/aibridged_internal_test.go | 4 ++++ cli/server_aibridge_internal_test.go | 2 ++ 2 files changed, 6 insertions(+) diff --git a/cli/aibridged_internal_test.go b/cli/aibridged_internal_test.go index d8d02c13146..59b2153c6f5 100644 --- a/cli/aibridged_internal_test.go +++ b/cli/aibridged_internal_test.go @@ -175,6 +175,8 @@ func TestBuildProviders(t *testing.T) { cfg.LegacyBedrock.Region = serpent.String("us-west-2") cfg.LegacyBedrock.AccessKey = serpent.String("AKID") cfg.LegacyBedrock.AccessKeySecret = serpent.String("secret") + cfg.LegacyBedrock.Model = serpent.String("anthropic.claude-3-5-sonnet-20241022-v2:0") + cfg.LegacyBedrock.SmallFastModel = serpent.String("anthropic.claude-3-5-haiku-20241022-v1:0") providers, err := buildFromEnv(t, cfg) require.NoError(t, err) @@ -191,6 +193,8 @@ func TestBuildProviders(t *testing.T) { cfg.LegacyBedrock.Region = serpent.String("us-west-2") cfg.LegacyBedrock.AccessKey = serpent.String("AKID") cfg.LegacyBedrock.AccessKeySecret = serpent.String("secret") + cfg.LegacyBedrock.Model = serpent.String("anthropic.claude-3-5-sonnet-20241022-v2:0") + cfg.LegacyBedrock.SmallFastModel = serpent.String("anthropic.claude-3-5-haiku-20241022-v1:0") providers, err := buildFromEnv(t, cfg) require.NoError(t, err) diff --git a/cli/server_aibridge_internal_test.go b/cli/server_aibridge_internal_test.go index 7a3732cef77..781b642f938 100644 --- a/cli/server_aibridge_internal_test.go +++ b/cli/server_aibridge_internal_test.go @@ -727,6 +727,8 @@ func TestBuildProviderFromProtoSetsAPIDumpDir(t *testing.T) { Region: "us-east-1", AccessKey: "AKID", AccessKeySecret: "secret", + Model: "anthropic.claude-3-5-sonnet-20241022-v2:0", + SmallFastModel: "anthropic.claude-3-5-haiku-20241022-v1:0", }, }, expectedType: aibridge.ProviderAnthropic, From 5429e02e7b80097d30460d410519ff99a9870885 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 14 Jul 2026 18:20:58 +0000 Subject: [PATCH 6/6] refactor: use resolved-protocol method throughout the codebase --- aibridge/config/config.go | 4 ++-- aibridge/intercept/messages/base.go | 5 ++--- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/aibridge/config/config.go b/aibridge/config/config.go index a3ba08f470c..ee0c5fec8a6 100644 --- a/aibridge/config/config.go +++ b/aibridge/config/config.go @@ -82,8 +82,8 @@ func (c AWSBedrock) ResolvedProtocol() BedrockProtocol { // Validate verifies protocol-specific Bedrock configuration. func (c AWSBedrock) Validate() error { - switch c.Protocol { - case "", BedrockProtocolInvokeModel: + switch c.ResolvedProtocol() { + case BedrockProtocolInvokeModel: if c.Region == "" && c.BaseURL == "" { return xerrors.New("region or base url required") } diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index 20f581bb59b..1a9ea0dfb85 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -133,14 +133,13 @@ func (i *interceptionBase) CorrelatingToolCallID() *string { // isBedrockMantle reports whether the interception targets the Bedrock mantle // protocol. func (i *interceptionBase) isBedrockMantle() bool { - return i.bedrock != nil && i.bedrock.Cfg.Protocol == aibconfig.BedrockProtocolMantle + return i.bedrock != nil && i.bedrock.Cfg.ResolvedProtocol() == aibconfig.BedrockProtocolMantle } // isBedrockInvokeModel reports whether the interception targets the Bedrock // InvokeModel protocol. func (i *interceptionBase) isBedrockInvokeModel() bool { - return i.bedrock != nil && - (i.bedrock.Cfg.Protocol == "" || i.bedrock.Cfg.Protocol == aibconfig.BedrockProtocolInvokeModel) + return i.bedrock != nil && i.bedrock.Cfg.ResolvedProtocol() == aibconfig.BedrockProtocolInvokeModel } func (i *interceptionBase) Model() string {